Fix and speed up the fast axis-angle conversions

Summary:
GitHub issue #2002 (https://github.com/facebookresearch/pytorch3d/issues/2002)
points out that the near-pi branch of `matrix_to_axis_angle(..., fast=True)`
normalizes with `torch.norm(n)`, which reduces over the whole batch. That is
real, and there were two further problems next to it.

`matrix_to_axis_angle(..., fast=True)`:

- The axis was read from row 0 of `(R + I) / 2`, which is `n_x * n`. That is
  exactly zero whenever the axis is perpendicular to x, giving `nan` (for
  example for an exact rotation by pi about y), and is dominated by rounding
  when `abs(n_x)` is small. It is now read from
  `R + R^T - 2*cos(angle)*I == 2*(1 - cos(angle))*nnT`, taking the column whose
  diagonal entry is largest so that the multiplier is at least `1/sqrt(3)`.
  Symmetrizing also makes the identity exact at every angle rather than only at
  pi. The sign, which `nnT` does not determine, comes from `omegas`.
- `torch.norm(n)` becomes a per-row `torch.linalg.vector_norm`, which is the
  reported bug: with more than one near-pi rotation in a batch, every one of
  them was scaled wrongly.
- The branch threshold was `isclose(angle, pi)`. `omegas` is `2*sin(angle)`
  times the axis, so the other branch loses relative precision like
  `1/(pi - angle)`, and just outside `isclose` the float32 round-trip error
  reached 9.6e-4. The threshold is now `pi - 1e-2`, which leaves under 1% of
  uniformly random rotations on the more expensive branch.
- The `torch.isclose(angles, 0)` guard on `omegas` is removed. `torch.sinc(0)`
  is 1, so a zero angle was never a special case, and `torch.norm` has a zero
  rather than `nan` gradient at zero. The guard only zeroed the answer for
  angles below its `atol` of 1e-8, returning 0 instead of an exact 1e-9
  rotation, and cost 24% at a batch of 100k.

Worst-case float32 round-trip matrix error over random axes, before -> after:
exactly pi `nan` or 2.0 -> 3.6e-07; `pi - 1e-5` 2.0 -> 4.2e-07; `pi - 1e-4`
9.6e-04 -> 3.6e-07. The worst case at any angle is now 9.9e-06, at the branch
boundary, against 4.8e-07 for `fast=False`.

On the choice of `pi - 1e-2`: the new near-pi branch is accurate at any
threshold, so the threshold only decides how much of the batch takes the
slower branch, and the worst case is always the angle just below it. Measured
on one batch of 100k random float32 rotations on CPU, where `fast=False` took
11.4 ms, threshold against fraction of uniformly random rotations selected,
time, and worst-case round-trip error:

- `isclose(angle, pi)`, the previous behaviour: 0.003%, 8.6 ms, 1e-03
- `pi - 1e-3`: 0.1%, 8.9 ms, 9e-05
- `pi - 1e-2`, chosen: 0.7%, 8.9 ms, 1e-05
- `pi - 0.05`: 3.2%, 10.2 ms, 3e-06
- `3.0` radians: 9.0%, 17.6 ms, 9e-07

So `pi - 1e-2` buys two orders of magnitude of accuracy over the old threshold
for no measurable time, and it is the last threshold that is free; at `3.0`
radians `fast=True` would be slower than `fast=False`. The errors are sampled
over random axes, so they move in the last digit between runs.

`axis_angle_to_matrix(..., fast=True)` is Rodrigues' formula rewritten with
`cross_product_matrix^2 == outer_product - angle^2 * I` and
`(1 - cos(angle)) / angle^2 == sinc(angle / (2*pi))^2 / 2`. The first removes a
batched 3x3 matrix multiplication, which `bmm` serves poorly at that size, and
folds the leftover `-angle^2 * I` into the identity term as `cos(angle) * I`;
the second is defined at zero, so the `angles_sqrd == 0` special case goes
away. Values are unchanged to 8.9e-16.

`matrix_to_axis_angle` also used eight separate boolean mask indexes, each of
which re-runs `nonzero` over the whole batch and, on CUDA, synchronizes. The
cheap branch is now evaluated densely and the near-pi minority is selected
with a single `nonzero`.

Evaluating the near-pi branch densely too, which would remove the last
`nonzero` and so every data-dependent shape, was tried and rejected: against
this version it is 1.24x to 1.60x faster on CUDA at batches of 1k to 100k but
0.88x at 1M, and about 2x slower on CPU, and it does not survive
`torch.jit.script` without further work. It is the version to revisit if
export or `torch.compile` friendliness ever matters more than speed.

Speedups, before -> after:

- `matrix_to_axis_angle`: CUDA 1.88x at a batch of 1k, 1.66x at 100k, 1.47x at
  1M; CPU 1.26x at 1k, 2.01x at 100k, 1.71x at 1M.
- `axis_angle_to_matrix`: CUDA 1.26x at 1k, 1.81x at 100k, 5.49x at 1M; CPU
  unchanged.

Before this change `matrix_to_axis_angle(..., fast=True)` was slower than
`fast=False` on CUDA at a batch of 100k, 0.587 against 0.463 ms, because the
selection synchronized.

Rewriting `axis_angle_to_matrix` to compute all nine entries in one
`torch.stack`, as `quaternion_to_matrix` does, was also tried and is not
faster: 0.96x, 0.74x and 1.10x on CUDA at 1k, 100k and 1M, because the extra
kernel launches cost about what the saved intermediates gain.

Reviewed By: MichaelRamamonjisoa

Differential Revision: D115714860

fbshipit-source-id: cf19695f67bf2e2e6f8719419e9f902d5c58c309

pytorch3d/csrc/marching_cubes/marching_cubes_cpu.cpp

@@ -34,6 +34,18 @@ std::tuple<at::Tensor, at::Tensor, at::Tensor> MarchingCubesCpu(
   const int H = vol.size(1);
   const int W = vol.size(2);
 
+  // The edge-id hash packs two grid-point ids as v1_id * (W*H*D) + v2_id, whose
+  // maximum value is (W*H*D)^2 - 1. Guard against signed int64 overflow so an
+  // oversized volume fails loudly instead of silently producing colliding edge
+  // ids and a corrupted mesh. floor(sqrt(INT64_MAX)) == 3037000499 (~1448^3).
+  const int64_t num_grid_points = (int64_t)W * H * D;
+  TORCH_CHECK(
+      num_grid_points <= 3037000499LL,
+      "Volume too large for CPU marching cubes: W*H*D (",
+      num_grid_points,
+      ") exceeds 3037000499 (~1448^3), which would overflow the int64 edge-id "
+      "hash and corrupt the mesh.");
+
   // Create tensor accessors
   auto vol_a = vol.accessor<float, 3>();
   // edge_id_to_v maps from an edge id to a vertex position
@@ -66,19 +78,23 @@ std::tuple<at::Tensor, at::Tensor, at::Tensor> MarchingCubesCpu(
           tri.push_back(edge);
           ps.push_back(interp_points[e]);
 
-          // Check if the triangle face is degenerate. A triangle face
-          // is degenerate if any of the two verices share the same 3D position
-          if ((j + 1) % 3 == 0 && ps[0] != ps[1] && ps[1] != ps[2] &&
-              ps[2] != ps[0]) {
-            for (int k = 0; k < 3; k++) {
-              int64_t v = tri.at(k);
-              edge_id_to_v[v] = ps.at(k);
-              if (!uniq_edge_id.count(v)) {
-                uniq_edge_id[v] = verts.size();
-                verts.push_back(edge_id_to_v[v]);
+          if (ps.size() == 3) {
+            // Check if the triangle face is degenerate. A triangle face
+            // is degenerate if any of the two vertices share the same 3D
+            // position
+            if (ps[0] != ps[1] && ps[1] != ps[2] && ps[2] != ps[0]) {
+              for (int k = 0; k < 3; k++) {
+                int64_t v = tri.at(k);
+                edge_id_to_v[v] = ps.at(k);
+                if (!uniq_edge_id.count(v)) {
+                  uniq_edge_id[v] = verts.size();
+                  verts.push_back(edge_id_to_v[v]);
+                }
+                faces.push_back(uniq_edge_id[v]);
               }
-              faces.push_back(uniq_edge_id[v]);
             }
+            // Clear unconditionally - a rejected degenerate triangle must not
+            // corrupt the buffer for the next triangle in this cube.
             tri.clear();
             ps.clear();
           } // endif

pytorch3d/csrc/marching_cubes/marching_cubes_utils.h

@@ -126,8 +126,7 @@ struct Cube {
   // edge with an integer to address floating point precision issue.
   //
   // Args:
-  //    v1_id: global id of vertex 1
-  //    v2_id: global id of vertex 2
+  //    edge: edge (ID) whose two endpoint vertices are hashed
   //    W: width of the 3d grid
   //    H: height of the 3d grid
   //    D: depth of the 3d grid
@@ -135,11 +134,19 @@ struct Cube {
   // Returns:
   //    hashing for a pair of vertex ids
   //
-  int64_t HashVpair(const int edge, int W, int H, int D) {
+  int64_t HashVpair(const int edge, int64_t W, int64_t H, int64_t D) const {
     const int v1 = _EDGE_TO_VERTICES[edge][0];
     const int v2 = _EDGE_TO_VERTICES[edge][1];
-    const int v1_id = p[v1].x + p[v1].y * W + p[v1].z * W * H;
-    const int v2_id = p[v2].x + p[v2].y * W + p[v2].z * W * H;
-    return (int64_t)v1_id * (W + W * H + W * H * D) + (int64_t)v2_id;
+    // p[v].x/y/z hold integral corner coordinates, so casting to int64_t is
+    // exact. We cast before multiplication to avoid float32 precision loss at
+    // grids > 256^3.
+    const int64_t v1_id =
+        (int64_t)p[v1].x + (int64_t)p[v1].y * W + (int64_t)p[v1].z * W * H;
+    const int64_t v2_id =
+        (int64_t)p[v2].x + (int64_t)p[v2].y * W + (int64_t)p[v2].z * W * H;
+    // v_id max is (W*H*D - 1). A stride of W*H*D perfectly separates v1 and v2.
+    // Note: This hash fits in int64_t safely up to grids of ~1448^3.
+    const int64_t stride = W * H * D;
+    return v1_id * stride + v2_id;
   }
 };

pytorch3d/ops/marching_cubes.py

@@ -159,7 +159,6 @@ def marching_cubes_naive(
     """
     batched_verts, batched_faces = [], []
     D, H, W = vol_batch.shape[1:]
-
     # each edge is represented with its two endpoints (represented with global id)
     for i in range(len(vol_batch)):
         vol = vol_batch[i]
@@ -184,7 +183,7 @@ def marching_cubes_naive(
                     # triangle vertex IDs and positions
                     tri = []
                     ps = []
-                    for i, edge in enumerate(edge_indices):
+                    for edge in edge_indices:
                         interp_points[edge] = cube.vert_interp(thresh, edge, vol)
 
                         # Bind interpolated vertex with a global edge_id, which
@@ -196,20 +195,19 @@ def marching_cubes_naive(
                         )
                         tri.append(edge_id)
                         ps.append(interp_points[edge])
-                        # when the isolevel are the same as the edge endpoints, the interploated
-                        # vertices can share the same values, and lead to degenerate triangles.
-                        if (
-                            (i + 1) % 3 == 0
-                            and ps[0] != ps[1]
-                            and ps[1] != ps[2]
-                            and ps[2] != ps[0]
-                        ):
-                            for j, edge_id in enumerate(tri):
-                                edge_id_to_v[edge_id] = ps[j]
-                                if edge_id not in uniq_edge_id:
-                                    uniq_edge_id[edge_id] = len(verts)
-                                    verts.append(edge_id_to_v[edge_id])
-                            faces.append([uniq_edge_id[tri[j]] for j in range(3)])
+                        if len(ps) == 3:
+                            # when the isolevel are the same as the edge
+                            # endpoints, the interpolated vertices can share
+                            # the same values, and lead to degenerate
+                            # triangles.
+                            if ps[0] != ps[1] and ps[1] != ps[2] and ps[2] != ps[0]:
+                                for j, edge_id in enumerate(tri):
+                                    edge_id_to_v[edge_id] = ps[j]
+                                    if edge_id not in uniq_edge_id:
+                                        uniq_edge_id[edge_id] = len(verts)
+                                        verts.append(edge_id_to_v[edge_id])
+                                faces.append([uniq_edge_id[tri[j]] for j in range(3)])
+
                             tri = []
                             ps = []
 

pytorch3d/transforms/rotation_conversions.py

@@ -489,15 +489,18 @@ def axis_angle_to_matrix(axis_angle: torch.Tensor, fast: bool = False) -> torch.
     cross_product_matrix = torch.stack(
         [zeros, -rz, ry, rz, zeros, -rx, -ry, rx, zeros], dim=-1
     ).view(shape + (3,))
-    cross_product_matrix_sqrd = cross_product_matrix @ cross_product_matrix
+    outer_product = axis_angle.unsqueeze(-1) * axis_angle.unsqueeze(-2)
 
+    # This is the Rodrigues formula written with the identities
+    #     cross_product_matrix^2 == outer_product - angle^2 * I
+    #     (1 - cos(angle)) / angle^2 == sinc(angle / (2*pi))^2 / 2
+    # which save a matrix multiplication and, as sinc is defined at zero,
+    # remove the need to special-case a zero angle.
     identity = torch.eye(3, dtype=dtype, device=device)
-    angles_sqrd = angles * angles
-    angles_sqrd = torch.where(angles_sqrd == 0, 1, angles_sqrd)
     return (
-        identity.expand(cross_product_matrix.shape)
+        torch.cos(angles) * identity
         + torch.sinc(angles / torch.pi) * cross_product_matrix
-        + ((1 - torch.cos(angles)) / angles_sqrd) * cross_product_matrix_sqrd
+        + 0.5 * torch.sinc(angles / (2 * torch.pi)) ** 2 * outer_product
     )
 
 
@@ -536,24 +539,52 @@ def matrix_to_axis_angle(matrix: torch.Tensor, fast: bool = False) -> torch.Tens
     traces = torch.diagonal(matrix, dim1=-2, dim2=-1).sum(-1).unsqueeze(-1)
     angles = torch.atan2(norms, traces - 1)
 
-    zeros = torch.zeros(3, dtype=matrix.dtype, device=matrix.device)
-    omegas = torch.where(torch.isclose(angles, torch.zeros_like(angles)), zeros, omegas)
+    # omegas is 2*sin(angle) times the axis, so the relative precision of the
+    # axis it gives decays like 1/(pi - angle). Above the following angle the
+    # axis is taken from the symmetric part of the matrix instead. Switching
+    # over 0.01 radians before pi keeps the error close to that of the slow
+    # implementation, while leaving under 1% of uniformly random rotations on
+    # the more expensive branch.
+    near_pi = angles > torch.pi - 1e-2
 
-    near_pi = angles.isclose(angles.new_full((1,), torch.pi)).squeeze(-1)
+    # An angle of zero needs no special case: torch.sinc(0) is 1, and omegas
+    # is then zero, which is the right answer. torch.sinc is instead zero at
+    # an angle of pi, so those entries are given a harmless denominator here
+    # and overwritten below.
+    sincs = torch.sinc(angles / torch.pi)
+    axis_angles = omegas * (0.5 / torch.where(near_pi, 1.0, sincs))
 
-    axis_angles = torch.empty_like(omegas)
-    axis_angles[~near_pi] = (
-        0.5 * omegas[~near_pi] / torch.sinc(angles[~near_pi] / torch.pi)
+    # Rotations by nearly pi are a small minority, so it is worth selecting
+    # them rather than evaluating the following densely. The selection index is
+    # computed once, because each boolean mask index would recompute it (and,
+    # on the GPU, synchronize).
+    flat_axis_angles = axis_angles.reshape(-1, 3)
+    index = near_pi.reshape(-1).nonzero().squeeze(-1)
+    near_pi_matrix = matrix.reshape(-1, 3, 3).index_select(0, index)
+    near_pi_omegas = omegas.reshape(-1, 3).index_select(0, index)
+
+    # this derives from: R + R^T - 2*cos(angle)*I = 2*(1 - cos(angle))*nnT,
+    # together with trace(R) - 1 == 2*cos(angle)
+    double_cosines = traces.reshape(-1, 1, 1).index_select(0, index) - 1
+    nnT = (
+        near_pi_matrix
+        + near_pi_matrix.transpose(-1, -2)
+        - double_cosines * torch.eye(3, dtype=matrix.dtype, device=matrix.device)
     )
+    # Every column of nnT is a multiple of n. Taking the one whose diagonal
+    # entry is largest means the multiplier has absolute value at least
+    # 1/sqrt(3), so the column is never degenerate.
+    largest = torch.diagonal(nnT, dim1=-2, dim2=-1).argmax(dim=-1)
+    n = nnT.gather(-1, largest.reshape(-1, 1, 1).expand(-1, 3, 1)).squeeze(-1)
+    n = n / torch.linalg.vector_norm(n, dim=-1, keepdim=True)
+    # nnT determines n only up to sign, which omegas resolves. (At exactly pi
+    # the sign is arbitrary, and there omegas is zero.)
+    n = torch.where((n * near_pi_omegas).sum(-1, keepdim=True) < 0, -n, n)
 
-    # this derives from: nnT = (R + 1) / 2
-    n = 0.5 * (
-        matrix[near_pi][..., 0, :]
-        + torch.eye(1, 3, dtype=matrix.dtype, device=matrix.device)
+    near_pi_angles = angles.reshape(-1, 1).index_select(0, index)
+    return torch.index_copy(flat_axis_angles, 0, index, near_pi_angles * n).reshape(
+        axis_angles.shape
     )
-    axis_angles[near_pi] = angles[near_pi] * n / torch.norm(n)
-
-    return axis_angles
 
 
 def axis_angle_to_quaternion(axis_angle: torch.Tensor) -> torch.Tensor:

tests/test_marching_cubes.py

@@ -780,6 +780,79 @@ class TestMarchingCubes(TestCaseMixin, unittest.TestCase):
         self.assertClose(faces[0], expected_faces)
         self.assertTrue(verts[0].ge(-1).all() and verts[0].le(1).all())
 
+    def test_degenerate_triangle_keeps_later_faces(self):
+        # A cube whose first triangle is degenerate must not suppress the
+        # remaining triangles of the same cube. Here the isolevel coincides
+        # with the value of the outside corners, so every interpolated point
+        # snaps onto a corner and the cube's 4 candidate triangles collapse
+        # onto the 4 outside corners: triangles 1 and 4 become degenerate
+        # while triangles 2 and 3 stay valid.
+        volume_data = torch.ones(1, 2, 2, 2)  # (B, W, H, D)
+        volume_data[0, 1, 0, 0] = 0
+        volume_data[0, 1, 0, 1] = 0
+        volume_data[0, 0, 1, 1] = 0
+        volume_data[0, 1, 1, 1] = 0
+        volume_data = volume_data.permute(0, 3, 2, 1)  # (B, D, H, W)
+
+        # The four inside corners are separated from the four outside corners
+        # (1, 1, 0), (0, 0, 1), (0, 1, 0) and (0, 0, 0), which the surface
+        # passes exactly through, giving a quad made of two triangles.
+        expected_verts = torch.tensor(
+            [
+                [1.0, 1.0, 0.0],
+                [0.0, 0.0, 1.0],
+                [0.0, 1.0, 0.0],
+                [0.0, 0.0, 0.0],
+            ]
+        )
+        expected_faces = torch.tensor([[0, 1, 2], [0, 3, 1]])
+
+        verts, faces = marching_cubes_naive(volume_data, 1, return_local_coords=False)
+        self.assertClose(verts[0], expected_verts)
+        self.assertClose(faces[0], expected_faces)
+
+        verts, faces = marching_cubes(volume_data, 1, return_local_coords=False)
+        self.assertClose(verts[0], expected_verts)
+        self.assertClose(faces[0], expected_faces)
+
+    def test_large_grid_edge_ids(self):
+        # The C++ implementation identifies a vertex by hashing the pair of
+        # grid-point ids of the edge it lies on. Grid-point ids run up to
+        # W * H * D - 1, so once the grid exceeds 2 ** 24 points the ids must
+        # not be computed in float32 or distinct edges collide and their
+        # vertices get merged. Use a long, thin volume to cross that bound
+        # with a volume small enough to allocate (~67MB).
+        W, H, D = 4_200_000, 2, 2
+        # Isolated points on the (y=1, z=1) row, whose ids are x + W + W * H
+        # and therefore the largest in the grid. Spacing 3 keeps the blobs
+        # from sharing a cube while cycling through every id residue.
+        first_x = 2**24 - W * H - W
+        xs = [first_x + 3 * i for i in range(16)]
+        self.assertLess(xs[-1] + 1, W)
+
+        vol = torch.ones(1, D, H, W)
+        for x in xs:
+            vol[0, 1, 1, x] = 0.0
+        verts, faces = marching_cubes(vol, 0.5, return_local_coords=False)
+
+        # Each isolated point cuts the four grid edges leading away from it,
+        # so every blob contributes 4 distinct vertices and 2 faces.
+        expected_verts = set()
+        for x in xs:
+            expected_verts.update(
+                [
+                    (x - 0.5, 1.0, 1.0),
+                    (x + 0.5, 1.0, 1.0),
+                    (float(x), 0.5, 1.0),
+                    (float(x), 1.0, 0.5),
+                ]
+            )
+        self.assertEqual(
+            {tuple(v) for v in verts[0].tolist()},
+            expected_verts,
+        )
+        self.assertEqual(faces[0].shape[0], 2 * len(xs))
+
     def test_sphere(self):
         # (B, W, H, D)
         volume = torch.Tensor(

tests/test_rotation_conversions.py

@@ -228,6 +228,165 @@ class TestRotationConversion(TestCaseMixin, unittest.TestCase):
         self.assertClose(data, axis_angle_to_matrix(euler_angles_fast))
         self.assertClose(data, axis_angle_to_matrix(euler_angles, fast=True))
 
+    def test_axis_angle_by_pi(self):
+        """Rotations by exactly pi are recovered, and don't depend on the batch."""
+        options = [0.0, -1.0, 1.0]
+        axes = torch.nn.functional.normalize(
+            torch.stack(
+                [
+                    torch.tensor(vec, dtype=torch.float64)
+                    for vec in itertools.islice(  # exclude [0, 0, 0]
+                        itertools.product(options, options, options), 1, None
+                    )
+                ]
+            ),
+            dim=-1,
+        )
+        # Rotation by pi around unit vector x is given by 2 x x^T - Id.
+        R = 2 * torch.matmul(axes[..., None], axes[..., None, :]) - torch.eye(
+            3, dtype=torch.float64
+        )
+        for fast in [False, True]:
+            axis_angles = matrix_to_axis_angle(R, fast=fast)
+            self.assertClose(
+                axis_angles.norm(dim=-1),
+                torch.full(axes.shape[:1], math.pi, dtype=torch.float64),
+            )
+            self.assertClose(axis_angle_to_matrix(axis_angles), R)
+            # The axis of one rotation cannot depend on the others.
+            singly = torch.stack(
+                [matrix_to_axis_angle(r, fast=fast) for r in R.unbind()]
+            )
+            self.assertClose(singly, axis_angles)
+
+    def test_axis_angle_near_pi(self):
+        """The fast path stays accurate as the angle approaches pi."""
+        for dtype, atol in [(torch.float64, 1e-9), (torch.float32, 1e-4)]:
+            axes = torch.nn.functional.normalize(
+                torch.randn(50, 3, dtype=dtype), dim=-1
+            )
+            for angle in [
+                math.pi,
+                math.pi - 1e-7,
+                math.pi - 1e-5,
+                math.pi - 1e-3,
+                math.pi - 1e-2,
+                math.pi - 0.1,
+                2.0,
+            ]:
+                data = axes * angle
+                R = axis_angle_to_matrix(data)
+                axis_angles = matrix_to_axis_angle(R, fast=True)
+                self.assertClose(axis_angle_to_matrix(axis_angles), R, atol=atol)
+                self.assertClose(
+                    axis_angles.norm(dim=-1),
+                    torch.full(axes.shape[:1], angle, dtype=dtype),
+                    atol=atol,
+                )
+                if angle < math.pi - 1e-2:
+                    # Closer to pi than this, the sign of the axis is not
+                    # determined to float32 precision (and both signs describe
+                    # the same rotation to within the tolerance above).
+                    self.assertClose(axis_angles, data, atol=atol)
+
+    def test_axis_angle_to_matrix_fast(self):
+        """The Rodrigues implementation agrees with the quaternion one."""
+        data = torch.randn(100, 3, dtype=torch.float64)
+        normalized = torch.nn.functional.normalize(data, dim=-1)
+        data[:5] = 0.0
+        data[5:10] = normalized[5:10] * math.pi
+        data[10:15] = normalized[10:15] * 1e-12
+        self.assertClose(
+            axis_angle_to_matrix(data, fast=True),
+            axis_angle_to_matrix(data),
+            atol=1e-14,
+        )
+
+        # A zero angle is not a special case, but must still have gradients.
+        zeros = torch.zeros(1, 3, dtype=torch.float64, requires_grad=True)
+        matrices = axis_angle_to_matrix(zeros, fast=True)
+        self.assertClose(matrices, torch.eye(3, dtype=torch.float64)[None])
+        (grad,) = torch.autograd.grad(matrices.sum(), zeros)
+        self.assertTrue(torch.isfinite(grad).all())
+
+    def test_axis_angle_tiny(self):
+        """Tiny angles are preserved, not flattened to zero, and have grads."""
+        axis = torch.nn.functional.normalize(
+            torch.tensor([[0.3, -0.5, 0.81]], dtype=torch.float64), dim=-1
+        )
+        for angle in [1e-3, 1e-8, 1e-9, 1e-12]:
+            data = axis * angle
+            R = axis_angle_to_matrix(data)
+            for fast in [False, True]:
+                # atol must be 0 here: the default would accept zero.
+                self.assertClose(matrix_to_axis_angle(R, fast=fast), data, atol=0)
+
+        identity = torch.eye(3, dtype=torch.float64)
+        for fast in [False, True]:
+            R = identity.clone().requires_grad_(True)
+            axis_angles = matrix_to_axis_angle(R, fast=fast)
+            self.assertClose(axis_angles, torch.zeros(3, dtype=torch.float64))
+            (grad,) = torch.autograd.grad(axis_angles.sum(), R)
+            self.assertTrue(torch.isfinite(grad).all())
+
+    def test_axis_angle_batch_independence(self):
+        """Every matrix in a batch is converted independently of the others."""
+        angles = torch.tensor(
+            [math.pi, math.pi, math.pi - 1e-9, math.pi - 1e-3, 3.0, 1.0, 0.5, 0.0],
+            dtype=torch.float64,
+        )
+        axes = torch.nn.functional.normalize(
+            torch.randn(angles.shape[0], 3, dtype=torch.float64), dim=-1
+        )
+        R = axis_angle_to_matrix(axes * angles[:, None])
+        for fast in [False, True]:
+            batched = matrix_to_axis_angle(R, fast=fast)
+            singly = torch.stack(
+                [matrix_to_axis_angle(r, fast=fast) for r in R.unbind()]
+            )
+            self.assertClose(batched, singly)
+
+    def test_axis_angle_shapes(self):
+        """Arbitrary leading dimensions, including none and zero."""
+        R = random_rotations(24, dtype=torch.float64).reshape(2, 3, 4, 3, 3)
+        # A rotation by pi, so that the near-pi branch is used too.
+        axis = torch.tensor([0.0, 0.6, -0.8], dtype=torch.float64)
+        R[1, 2, 0] = 2 * torch.outer(axis, axis) - torch.eye(3, dtype=torch.float64)
+        for fast in [False, True]:
+            axis_angles = matrix_to_axis_angle(R, fast=fast)
+            self.assertEqual(axis_angles.shape, (2, 3, 4, 3))
+            self.assertClose(axis_angle_to_matrix(axis_angles), R)
+
+            self.assertEqual(matrix_to_axis_angle(R[0, 0, 0], fast=fast).shape, (3,))
+            empty = torch.zeros(0, 3, 3, dtype=torch.float64)
+            self.assertEqual(matrix_to_axis_angle(empty, fast=fast).shape, (0, 3))
+
+    def test_axis_angle_near_pi_grad(self):
+        """Round tripping is the identity, so its derivative is too."""
+        axis = torch.nn.functional.normalize(
+            torch.tensor([0.3, -0.5, 0.81], dtype=torch.float64), dim=-1
+        )
+        eye = torch.eye(3, dtype=torch.float64)
+        for angle in [0.5, 2.0, math.pi - 0.02, math.pi - 1e-4]:
+            data = axis * angle
+            for fast in [False, True]:
+                jacobian = torch.autograd.functional.jacobian(
+                    lambda x, fast=fast: matrix_to_axis_angle(
+                        axis_angle_to_matrix(x), fast=fast
+                    ),
+                    data,
+                )
+                self.assertClose(jacobian, eye, atol=1e-8)
+
+        # At exactly pi the axis, and so the derivative, is ambiguous, but the
+        # result must still be usable.
+        R = (2 * torch.outer(axis, axis) - eye).requires_grad_(True)
+        for fast in [False, True]:
+            (grad,) = torch.autograd.grad(
+                matrix_to_axis_angle(R, fast=fast).sum(), R, retain_graph=True
+            )
+            self.assertTrue(torch.isfinite(grad).all())
+
     def test_quaternion_application(self):
         """Applying a quaternion is the same as applying the matrix."""
         quaternions = random_quaternions(3, torch.float64)