docstr: Checks that for an identity rotation `R=I`, the so3_rotation_angle returns
  non-finite gradients when `cos_bound=None` and finite gradients
  for `cos_bound > 0.0`.
docstr_tok: ["Checks","that","for","an","identity","rotation","`","R=I","`",",","the","so3_rotation_angle","returns","non-finite","gradients","when","`","cos_bound=None","`","and","finite","gradients","for","`","cos_bound",">","0.0","`","."]
code: def test_so3_cos_bound(self, batch_size: int = 100):
      """
      Checks that for an identity rotation `R=I`, the so3_rotation_angle returns
      non-finite gradients when `cos_bound=None` and finite gradients
      for `cos_bound > 0.0`.
      """
      # generate random rotations with a tiny angle to generate cases
      # with the gradient singularity
      device = torch.device("cuda:0")
      identity = torch.eye(3, device=device)
      rot180 = identity * torch.tensor(
          [[1.0, -1.0, -1.0]], device=device
      )
      r = [identity, rot180]
      r.extend(
          [
              torch.linalg.qr(
                  identity + torch.randn_like(identity) * 1e-4
              )[0]
              for _ in range(batch_size - 2)
          ]
      )
      r = torch.stack(r)
      r.requires_grad = True
      for is_grad_finite in (True, False):
          # clear the gradients and decide the cos_bound:
          #     for is_grad_finite we run so3_rotation_angle with cos_bound
          #     set to a small float, otherwise we set to 0.0
          r.grad = None
          cos_bound = 1e-4 if is_grad_finite else 0.0
          # compute the angles of r
          angles = so3_rotation_angle(r, cos_bound=cos_bound)
          # tests whether all outputs are finite in both cases
          self.assertTrue(torch.isfinite(angles).all())
          # compute the gradients
          loss = angles.sum()
          loss.backward()
          # tests whether the gradient is not None for both cases
          self.assertIsNotNone(r.grad)
          if is_grad_finite:
              # all grad values have to be finite
              self.assertTrue(torch.isfinite(r.grad).all())
  
imports: import math
  import unittest
  from distutils.version import LooseVersion
  import numpy
  import torch
  from pytorch3d.transforms.so3 import hat, so3_exp_map, so3_log_map, so3_relative_angle, so3_rotation_angle
  from .common_testing import TestCaseMixin
label_desc: Use image node_id 14 for calling the TestSO3 obj's underlying member method code with example usage: obj.test_so3_cos_bound(batch_size) without return types