diff --git a/gnm/shape/gnm_common.py b/gnm/shape/gnm_common.py
index f0b9ef25..34251cd7 100644
--- a/gnm/shape/gnm_common.py
+++ b/gnm/shape/gnm_common.py
@@ -234,21 +234,21 @@ def axis_angle_to_rotation_matrix(
   return matrix
 
 
-def joint_transforms_world(
+def joint_transforms_world_from_rotation_matrices(
     joints: enpt.FloatArray['... J 3'],
-    rotations: enpt.FloatArray['... J 3'],
+    rotation_matrices: enpt.FloatArray['... J 3 3'],
     translation: enpt.FloatArray['... 3'],
     joint_parent_indices: Sequence[int],
 ) -> enpt.FloatArray['... J 4 4']:
-  """Computes the world-space transformation matrices for each joint.
+  """Computes world-space joint transforms from rotation matrices.
 
   This function traverses the joint hierarchy and applies forward kinematics
   to compute the global/world space transforms of each joint, combining
-  joint rotations and joint translations.
+  joint rotation matrices and joint translations.
 
   Args:
     joints: The bind-pose joints positions. Shape: `(..., J, 3)`.
-    rotations: The joint rotations as axis-angle vectors. Shape: `(..., J, 3)`.
+    rotation_matrices: The joint rotation matrices. Shape: `(..., J, 3, 3)`.
     translation: The translation vector for the root joint. Shape: `(..., 3)`.
     joint_parent_indices: A sequence of parent indices for each joint. The root
       joint must have itself or a sentinel as parent (usually index 0 has parent
@@ -269,11 +269,6 @@ def joint_transforms_world(
   )
   joint_parents = take(joints, parent_indices_except_root, axis=-2, xnp=xnp)
 
-  # Since we cannot do in-place assignments, we construct the local
-  # transformation matrices by combining rotation matrices and translation
-  # vectors.
-  local_rotations = axis_angle_to_rotation_matrix(rotations)
-
   # Construct local translations.
   root_translation = joints[..., 0, :] + translation
   non_root_translations = joints[..., 1:, :] - joint_parents
@@ -284,7 +279,7 @@ def joint_transforms_world(
   # Combine rotation and translation.
   # local_transforms: (..., J, 3, 4)
   local_transforms_3x4 = xnp.concatenate(
-      [local_rotations, local_translations[..., None]], axis=-1
+      [rotation_matrices, local_translations[..., None]], axis=-1
   )
 
   # Append [0, 0, 0, 1] row.
@@ -316,49 +311,71 @@ def joint_transforms_world(
   return xnp.stack(world_transforms_list, axis=-3)
 
 
-def linear_blend_skinning(
-    local_vertices: enpt.FloatArray['... V 3'],
-    joint_positions: enpt.FloatArray['... J 3'],
+def joint_transforms_world(
+    joints: enpt.FloatArray['... J 3'],
     rotations: enpt.FloatArray['... J 3'],
     translation: enpt.FloatArray['... 3'],
-    skinning_weights: enpt.FloatArray['J V'],
     joint_parent_indices: Sequence[int],
-) -> enpt.FloatArray['... V 3']:
-  """Poses the GNM vertices using Linear Blend Skinning (LBS).
+) -> enpt.FloatArray['... J 4 4']:
+  """Computes the world-space transformation matrices for each joint.
 
-  Linear Blend Skinning computes posed vertex positions by taking a weighted
-  sum of the transformations applied to each joint, where the weights are
-  defined per vertex and per joint.
+  This function traverses the joint hierarchy and applies forward kinematics
+  to compute the global/world space transforms of each joint, combining
+  joint rotations and joint translations.
 
   Args:
-    local_vertices: The vertex positions in the local character space, with
-      shape `(..., V, 3)`.
+    joints: The bind-pose joints positions. Shape: `(..., J, 3)`.
+    rotations: The joint rotations as axis-angle vectors. Shape: `(..., J, 3)`.
+    translation: The translation vector for the root joint. Shape: `(..., 3)`.
+    joint_parent_indices: A sequence of parent indices for each joint. The root
+      joint must have itself or a sentinel as parent (usually index 0 has parent
+      0).
+
+  Returns:
+    A batch of 4x4 transformation matrices in world space for each joint.
+    Shape: `(..., J, 4, 4)`.
+  """
+  rotation_matrices = axis_angle_to_rotation_matrix(rotations)
+  return joint_transforms_world_from_rotation_matrices(
+      joints=joints,
+      rotation_matrices=rotation_matrices,
+      translation=translation,
+      joint_parent_indices=joint_parent_indices,
+  )
+
+
+def apply_linear_blend_skinning_with_transforms(
+    local_vertices: enpt.FloatArray['... V 3'],
+    joint_positions: enpt.FloatArray['... J 3'],
+    joint_transforms_world_local: enpt.FloatArray['... J 4 4'],
+    skinning_weights: enpt.FloatArray['J V'],
+    xnp: enp.NpModule | None = None,
+) -> enpt.FloatArray['... V 3']:
+  """Applies Linear Blend Skinning using precomputed world joint transforms.
+
+  Args:
+    local_vertices: The vertex positions in local character space, with shape
+      `(..., V, 3)`.
     joint_positions: The joint positions in bind pose, with shape `(..., J, 3)`.
-    rotations: Joint rotations as axis-angle vectors, with shape `(..., J, 3)`.
-    translation: Translation vector for the root joint, with shape: `(..., 3)`.
-    skinning_weights: Skinning weights defining joint influence per vertex with
+    joint_transforms_world_local: The local-to-world 4x4 transformation matrices
+      of each joint after posing, with shape `(..., J, 4, 4)`.
+    skinning_weights: Skinning weights defining joint influence per vertex, with
       shape `(J, V)`.
-    joint_parent_indices: Sequence of parent indices for each joint.
+    xnp: Optional array module (e.g. `numpy`, `jax.numpy`, `torch`) to use. If
+      None, it is inferred from `local_vertices`.
 
   Returns:
-    The posed skinned vertices with shape `(..., V, 3)`.
+    The posed skinned vertices in world space, with shape `(..., V, 3)`.
   """
-  xnp = enp.get_np_module(local_vertices)
-
-  # The local-to-world transforms of each joint, after posing, (..., J, 4, 4).
-  joint_transforms_world_local = joint_transforms_world(
-      joint_positions, rotations, translation, joint_parent_indices
-  )
+  if xnp is None:
+    xnp = enp.get_np_module(local_vertices)
 
-  # Extract 3x3 rotation and translation from joint transforms.
   r_world = joint_transforms_world_local[..., :3, :3]
   t_world = joint_transforms_world_local[..., :3, 3]
 
-  # Compute bind-to-world translation: T_bind = T_world - R_world @ joints
   deltas = (r_world @ joint_positions[..., :, None])[..., 0]
   t_bind_to_world = t_world - deltas
 
-  # Compute per-vertex 3x3 rotations and 3D translations via skinning weights.
   per_vertex_r = xnp.einsum(
       'jv,...jmn->...vmn', skinning_weights, r_world, **_einsum_kwargs(xnp)
   )
@@ -369,7 +386,6 @@ def linear_blend_skinning(
       **_einsum_kwargs(xnp),
   )
 
-  # Apply affine transform: V_skinned = R_v @ V_local + T_v
   vertices_skinned = (per_vertex_r @ local_vertices[..., :, None])[
       ..., 0
   ] + per_vertex_t
@@ -377,6 +393,87 @@ def linear_blend_skinning(
   return vertices_skinned
 
 
+def linear_blend_skinning_from_rotation_matrices(
+    local_vertices: enpt.FloatArray['... V 3'],
+    joint_positions: enpt.FloatArray['... J 3'],
+    rotation_matrices: enpt.FloatArray['... J 3 3'],
+    translation: enpt.FloatArray['... 3'],
+    skinning_weights: enpt.FloatArray['J V'],
+    joint_parent_indices: Sequence[int],
+) -> enpt.FloatArray['... V 3']:
+  """Poses GNM vertices using Linear Blend Skinning from rotation matrices.
+
+  Linear Blend Skinning computes posed vertex positions by taking a weighted
+  sum of the transformations applied to each joint, where the weights are
+  defined per vertex and per joint.
+
+  Args:
+    local_vertices: The vertex positions in the local character space, with
+      shape `(..., V, 3)`.
+    joint_positions: The joint positions in bind pose, with shape `(..., J, 3)`.
+    rotation_matrices: Joint rotation matrices, with shape `(..., J, 3, 3)`.
+    translation: Translation vector for the root joint, with shape: `(..., 3)`.
+    skinning_weights: Skinning weights defining joint influence per vertex with
+      shape `(J, V)`.
+    joint_parent_indices: Sequence of parent indices for each joint.
+
+  Returns:
+    The posed skinned vertices with shape `(..., V, 3)`.
+  """
+  xnp = enp.get_np_module(local_vertices)
+
+  # The local-to-world transforms of each joint, after posing, (..., J, 4, 4).
+  joint_transforms_world_local = joint_transforms_world_from_rotation_matrices(
+      joint_positions, rotation_matrices, translation, joint_parent_indices
+  )
+
+  return apply_linear_blend_skinning_with_transforms(
+      local_vertices=local_vertices,
+      joint_positions=joint_positions,
+      joint_transforms_world_local=joint_transforms_world_local,
+      skinning_weights=skinning_weights,
+      xnp=xnp,
+  )
+
+
+def linear_blend_skinning(
+    local_vertices: enpt.FloatArray['... V 3'],
+    joint_positions: enpt.FloatArray['... J 3'],
+    rotations: enpt.FloatArray['... J 3'],
+    translation: enpt.FloatArray['... 3'],
+    skinning_weights: enpt.FloatArray['J V'],
+    joint_parent_indices: Sequence[int],
+) -> enpt.FloatArray['... V 3']:
+  """Poses the GNM vertices using Linear Blend Skinning (LBS).
+
+  Linear Blend Skinning computes posed vertex positions by taking a weighted
+  sum of the transformations applied to each joint, where the weights are
+  defined per vertex and per joint.
+
+  Args:
+    local_vertices: The vertex positions in the local character space, with
+      shape `(..., V, 3)`.
+    joint_positions: The joint positions in bind pose, with shape `(..., J, 3)`.
+    rotations: Joint rotations as axis-angle vectors, with shape `(..., J, 3)`.
+    translation: Translation vector for the root joint, with shape: `(..., 3)`.
+    skinning_weights: Skinning weights defining joint influence per vertex with
+      shape `(J, V)`.
+    joint_parent_indices: Sequence of parent indices for each joint.
+
+  Returns:
+    The posed skinned vertices with shape `(..., V, 3)`.
+  """
+  rotation_matrices = axis_angle_to_rotation_matrix(rotations)
+  return linear_blend_skinning_from_rotation_matrices(
+      local_vertices=local_vertices,
+      joint_positions=joint_positions,
+      rotation_matrices=rotation_matrices,
+      translation=translation,
+      skinning_weights=skinning_weights,
+      joint_parent_indices=joint_parent_indices,
+  )
+
+
 def vertex_positions_bind_pose(
     identity: enpt.FloatArray['... I'] | None,
     expression: enpt.FloatArray['... E'] | None,
@@ -401,7 +498,12 @@ def vertex_positions_bind_pose(
     The template vertices with shape and expression offsets applied. Shape:
     `(..., V, 3)`.
   """
-  xnp = enp.get_np_module(template_vertex_positions)
+  reference_array = (
+      identity
+      if identity is not None
+      else (expression if expression is not None else template_vertex_positions)
+  )
+  xnp = enp.get_np_module(reference_array)
 
   # Apply linear identity and expression bases.
   identity_deltas = 0.0
@@ -434,7 +536,10 @@ def joint_positions_bind_pose(
   Returns:
     The template joints with shape offsets applied. Shape: `(..., J, 3)`.
   """
-  xnp = enp.get_np_module(template_joint_positions)
+  reference_array = (
+      identity if identity is not None else template_joint_positions
+  )
+  xnp = enp.get_np_module(reference_array)
   deltas = 0.0
   if identity is not None:
     deltas = _apply_linear_basis(identity, joint_identity_basis, xnp)
@@ -442,20 +547,20 @@ def joint_positions_bind_pose(
   return template_joint_positions + deltas
 
 
-def compute_pose_correctives(
-    rotations: enpt.FloatArray['... J 3'] | None,
+def compute_pose_correctives_from_rotation_matrices(
+    rotation_matrices: enpt.FloatArray['... J 3 3'] | None,
     pose_correctives_regressor: enpt.FloatArray['F F2'] | None,
     template_vertex_positions: enpt.FloatArray['V 3'],
     num_joints: int,
     num_vertices: int,
 ) -> enpt.FloatArray['... V 3']:
-  """Applies pose-dependent corrective shape offsets to vertices.
+  """Applies pose-dependent corrective vertex offsets from rotation matrices.
 
-  Pose correctives modify the mesh vertices based on joint rotation angles to
+  Pose correctives modify the mesh vertices based on joint rotation matrices to
   fix skinning artifacts (e.g. at elbows, knees).
 
   Args:
-    rotations: Joint rotations as axis-angle vectors, with shape `(..., J, 3)`.
+    rotation_matrices: Joint rotation matrices, with shape `(..., J, 3, 3)`.
     pose_correctives_regressor: The regressor mapping pose features to vertex
       offsets, with shape `(F, F2)`.
     template_vertex_positions: The base template vertices, used as reference for
@@ -466,26 +571,30 @@ def compute_pose_correctives(
   Returns:
     The corrective vertex offsets. Shape: `(..., V, 3)`.
   """
-  xnp = enp.get_np_module(template_vertex_positions)
-
-  if pose_correctives_regressor is None or rotations is None:
+  if pose_correctives_regressor is None or rotation_matrices is None:
+    reference_array = (
+        rotation_matrices
+        if rotation_matrices is not None
+        else template_vertex_positions
+    )
     return zeros_with_batch_dims(
-        rotations if rotations is not None else template_vertex_positions,
-        2,
+        reference_array,
+        3 if rotation_matrices is not None else 2,
         (num_vertices, 3),
-        dtype=template_vertex_positions.dtype,
+        dtype=reference_array.dtype,
     )
 
-  rotation_matrices = axis_angle_to_rotation_matrix(rotations)
-
-  # Construct eye matrices.
-  eye_matrix = eye(3, dtype=rotations.dtype, reference_array=rotations, xnp=xnp)
-  # Broadcast to matching shape.
-  eye_reshaped = xnp.reshape(eye_matrix, (1,) * (rotations.ndim - 1) + (3, 3))
-  pose_features = rotation_matrices - eye_reshaped
+  xnp = enp.get_np_module(rotation_matrices)
+  eye_matrix = eye(
+      3,
+      dtype=rotation_matrices.dtype,
+      reference_array=rotation_matrices,
+      xnp=xnp,
+  )
+  pose_features = rotation_matrices - eye_matrix
 
   pose_features = reshape_with_batch_dims(
-      pose_features, (num_joints * 9,), rotations, 2
+      pose_features, (num_joints * 9,), rotation_matrices, 3
   )
 
   pose_deltas = xnp.einsum(
@@ -494,7 +603,47 @@ def compute_pose_correctives(
       pose_correctives_regressor,
       **_einsum_kwargs(xnp),
   )
-  return reshape_with_batch_dims(pose_deltas, (num_vertices, 3), rotations, 2)
+  return reshape_with_batch_dims(
+      pose_deltas, (num_vertices, 3), rotation_matrices, 3
+  )
+
+
+def compute_pose_correctives(
+    rotations: enpt.FloatArray['... J 3'] | None,
+    pose_correctives_regressor: enpt.FloatArray['F F2'] | None,
+    template_vertex_positions: enpt.FloatArray['V 3'],
+    num_joints: int,
+    num_vertices: int,
+) -> enpt.FloatArray['... V 3']:
+  """Applies pose-dependent corrective shape offsets to vertices.
+
+  Pose correctives modify the mesh vertices based on joint rotation angles to
+  fix skinning artifacts (e.g. at elbows, knees).
+
+  Args:
+    rotations: Joint rotations as axis-angle vectors, with shape `(..., J, 3)`.
+    pose_correctives_regressor: The regressor mapping pose features to vertex
+      offsets, with shape `(F, F2)`.
+    template_vertex_positions: The base template vertices, used as reference for
+      device and shape, with shape: `(V, 3)`.
+    num_joints: The total number of joints.
+    num_vertices: The total number of vertices.
+
+  Returns:
+    The corrective vertex offsets. Shape: `(..., V, 3)`.
+  """
+  rotation_matrices = (
+      axis_angle_to_rotation_matrix(rotations)
+      if rotations is not None
+      else None
+  )
+  return compute_pose_correctives_from_rotation_matrices(
+      rotation_matrices=rotation_matrices,
+      pose_correctives_regressor=pose_correctives_regressor,
+      template_vertex_positions=template_vertex_positions,
+      num_joints=num_joints,
+      num_vertices=num_vertices,
+  )
 
 
 def _graph_shape(
@@ -543,7 +692,12 @@ def _apply_linear_basis(
   batch_dims = params.shape[:-1]
 
   has_batch = batch_dims and (0 not in params.shape)
-  if has_batch and enp.lazy.is_np_xnp(xnp) and _is_constant_batch(params, xnp):
+  if (
+      has_batch
+      and enp.lazy.is_np(params)
+      and enp.lazy.is_np_xnp(xnp)
+      and _is_constant_batch(params, xnp)
+  ):
     single_param = params.reshape(-1, params.shape[-1])[0]
     deltas = xnp.einsum('i,ijk->jk', single_param, basis, **_einsum_kwargs(xnp))
     return xnp.broadcast_to(deltas, (*batch_dims, *deltas.shape))
diff --git a/gnm/shape/gnm_common_test.py b/gnm/shape/gnm_common_test.py
index 5046a426..10539033 100644
--- a/gnm/shape/gnm_common_test.py
+++ b/gnm/shape/gnm_common_test.py
@@ -23,6 +23,7 @@
 from absl.testing import parameterized
 from etils import enp
 from gnm.shape import gnm_common
+import jax
 import numpy as np
 
 # The array backends supported by the GNM math, as (name, xnp module) pairs.
@@ -212,6 +213,27 @@ def test_root_rotation_propagates_to_child(self, xnp: enp.NpModule):
         transforms[..., :3, 3][1], [0.0, 1.0, 0.0], atol=1e-5
     )
 
+  @parameterized.named_parameters(*_XNP_BACKENDS)
+  def test_from_rotation_matrices_matches_rotations(self, xnp: enp.NpModule):
+    rotations = xnp.asarray(
+        [[0.1, -0.2, 0.3], [0.4, 0.5, -0.6]], dtype=xnp.float32
+    )
+    rotation_matrices = gnm_common.axis_angle_to_rotation_matrix(rotations)
+    translation = xnp.asarray([1.0, 2.0, 3.0], dtype=xnp.float32)
+    joints = xnp.asarray(self.joints, dtype=xnp.float32)
+
+    transforms_rot = np.asarray(
+        gnm_common.joint_transforms_world(
+            joints, rotations, translation, self.parents
+        )
+    )
+    transforms_mat = np.asarray(
+        gnm_common.joint_transforms_world_from_rotation_matrices(
+            joints, rotation_matrices, translation, self.parents
+        )
+    )
+    np.testing.assert_allclose(transforms_mat, transforms_rot, atol=1e-5)
+
 
 class LinearBlendSkinningTest(parameterized.TestCase):
   """Tests for linear_blend_skinning."""
@@ -270,6 +292,135 @@ def test_preserves_batch_shape(self, xnp: enp.NpModule):
     )
     self.assertEqual(tuple(posed.shape), (4, 3, 3))
 
+  @parameterized.named_parameters(*_XNP_BACKENDS)
+  def test_from_rotation_matrices_matches_rotations(self, xnp: enp.NpModule):
+    rotations = xnp.asarray(
+        [[0.1, -0.2, 0.3], [0.4, 0.5, -0.6]], dtype=xnp.float32
+    )
+    rotation_matrices = gnm_common.axis_angle_to_rotation_matrix(rotations)
+    translation = xnp.asarray([1.0, 2.0, 3.0], dtype=xnp.float32)
+    vertices = xnp.asarray(self.vertices, dtype=xnp.float32)
+    joints = xnp.asarray(self.joints, dtype=xnp.float32)
+    weights = xnp.asarray(self.weights, dtype=xnp.float32)
+
+    posed_rot = np.asarray(
+        gnm_common.linear_blend_skinning(
+            vertices, joints, rotations, translation, weights, self.parents
+        )
+    )
+    posed_mat = np.asarray(
+        gnm_common.linear_blend_skinning_from_rotation_matrices(
+            vertices,
+            joints,
+            rotation_matrices,
+            translation,
+            weights,
+            self.parents,
+        )
+    )
+    np.testing.assert_allclose(posed_mat, posed_rot, atol=1e-5)
+
+  @parameterized.named_parameters(*_XNP_BACKENDS)
+  def test_apply_with_transforms_matches_from_rotation_matrices(
+      self, xnp: enp.NpModule
+  ) -> None:
+    rotations = xnp.asarray(
+        [[0.1, -0.2, 0.3], [0.4, 0.5, -0.6]], dtype=xnp.float32
+    )
+    rotation_matrices = gnm_common.axis_angle_to_rotation_matrix(rotations)
+    translation = xnp.asarray([1.0, 2.0, 3.0], dtype=xnp.float32)
+    vertices = xnp.asarray(self.vertices, dtype=xnp.float32)
+    joints = xnp.asarray(self.joints, dtype=xnp.float32)
+    weights = xnp.asarray(self.weights, dtype=xnp.float32)
+
+    transforms = gnm_common.joint_transforms_world_from_rotation_matrices(
+        joints, rotation_matrices, translation, self.parents
+    )
+    posed_with_transforms = np.asarray(
+        gnm_common.apply_linear_blend_skinning_with_transforms(
+            local_vertices=vertices,
+            joint_positions=joints,
+            joint_transforms_world_local=transforms,
+            skinning_weights=weights,
+        )
+    )
+    posed_from_mat = np.asarray(
+        gnm_common.linear_blend_skinning_from_rotation_matrices(
+            vertices,
+            joints,
+            rotation_matrices,
+            translation,
+            weights,
+            self.parents,
+        )
+    )
+    np.testing.assert_allclose(posed_with_transforms, posed_from_mat, atol=1e-5)
+
+  def test_jax_symbolic_batch_with_numpy_template_arrays(self) -> None:
+    (symbolic_b,) = jax.export.symbolic_shape('b')
+    identity_spec = jax.ShapeDtypeStruct((symbolic_b, 2), jax.numpy.float32)
+    expression_spec = jax.ShapeDtypeStruct((symbolic_b, 2), jax.numpy.float32)
+    rotation_matrices_spec = jax.ShapeDtypeStruct(
+        (symbolic_b, 2, 3, 3), jax.numpy.float32
+    )
+    translation_spec = jax.ShapeDtypeStruct((symbolic_b, 3), jax.numpy.float32)
+
+    template_vertices = np.asarray(self.vertices, dtype=np.float32)
+    template_joints = np.asarray(self.joints, dtype=np.float32)
+    vertex_identity_basis = np.ones((2, 3, 3), dtype=np.float32) * 0.01
+    joint_identity_basis = np.ones((2, 2, 3), dtype=np.float32) * 0.02
+    expression_basis = np.ones((2, 3, 3), dtype=np.float32) * 0.03
+    skinning_weights = np.asarray(self.weights, dtype=np.float32)
+    pose_correctives_regressor = (
+        np.arange(2 * 9 * 3 * 3).reshape(2 * 9, 3 * 3).astype(np.float32)
+        * 0.001
+    )
+
+    def _pipeline_fn(identity, expression, rotation_matrices, translation):
+      vertices_bind_pose = gnm_common.vertex_positions_bind_pose(
+          identity=identity,
+          expression=expression,
+          template_vertex_positions=template_vertices,
+          vertex_identity_basis=vertex_identity_basis,
+          expression_basis=expression_basis,
+      )
+      joints = gnm_common.joint_positions_bind_pose(
+          identity=identity,
+          template_joint_positions=template_joints,
+          joint_identity_basis=joint_identity_basis,
+      )
+      pose_correctives = (
+          gnm_common.compute_pose_correctives_from_rotation_matrices(
+              rotation_matrices=rotation_matrices,
+              pose_correctives_regressor=pose_correctives_regressor,
+              template_vertex_positions=template_vertices,
+              num_joints=2,
+              num_vertices=3,
+          )
+      )
+      joint_transforms = (
+          gnm_common.joint_transforms_world_from_rotation_matrices(
+              joints, rotation_matrices, translation, tuple(self.parents)
+          )
+      )
+      vertices_posed = gnm_common.apply_linear_blend_skinning_with_transforms(
+          local_vertices=vertices_bind_pose + pose_correctives,
+          joint_positions=joints,
+          joint_transforms_world_local=joint_transforms,
+          skinning_weights=skinning_weights,
+      )
+      return vertices_bind_pose, vertices_posed, joint_transforms
+
+    exported = jax.export.export(jax.jit(_pipeline_fn))(
+        identity_spec,
+        expression_spec,
+        rotation_matrices_spec,
+        translation_spec,
+    )
+    self.assertEqual(exported.out_avals[0].shape, (symbolic_b, 3, 3))
+    self.assertEqual(exported.out_avals[1].shape, (symbolic_b, 3, 3))
+    self.assertEqual(exported.out_avals[2].shape, (symbolic_b, 2, 4, 4))
+
 
 class BindPoseTest(parameterized.TestCase):
   """Tests for the bind-pose vertex and joint helpers."""
@@ -376,6 +527,56 @@ def test_preserves_batch_shape(self, xnp: enp.NpModule):
     )
     self.assertEqual(tuple(result.shape), (4, 3, 3))
 
+  @parameterized.named_parameters(*_XNP_BACKENDS)
+  def test_from_rotation_matrices_matches_rotations(self, xnp: enp.NpModule):
+    template = xnp.asarray(np.ones((3, 3)), dtype=xnp.float32)
+    regressor = xnp.asarray(
+        np.arange(2 * 9 * 3 * 3).reshape(2 * 9, 3 * 3) / 100.0,
+        dtype=xnp.float32,
+    )
+    rotations = xnp.asarray(
+        [[0.1, -0.2, 0.3], [0.4, 0.5, -0.6]], dtype=xnp.float32
+    )
+    rotation_matrices = gnm_common.axis_angle_to_rotation_matrix(rotations)
+
+    correctives_rot = np.asarray(
+        gnm_common.compute_pose_correctives(
+            rotations,
+            regressor,
+            template,
+            num_joints=2,
+            num_vertices=3,
+        )
+    )
+    correctives_mat = np.asarray(
+        gnm_common.compute_pose_correctives_from_rotation_matrices(
+            rotation_matrices,
+            regressor,
+            template,
+            num_joints=2,
+            num_vertices=3,
+        )
+    )
+    np.testing.assert_allclose(correctives_mat, correctives_rot, atol=1e-5)
+
+  @parameterized.named_parameters(*_XNP_BACKENDS)
+  def test_early_return_preserves_rotation_matrices_dtype(
+      self, xnp: enp.NpModule
+  ) -> None:
+    template = xnp.asarray(np.ones((3, 3)), dtype=xnp.float32)
+    rotation_matrices = xnp.asarray(
+        np.broadcast_to(np.eye(3), (2, 2, 3, 3)), dtype=xnp.float64
+    )
+    result = gnm_common.compute_pose_correctives_from_rotation_matrices(
+        rotation_matrices=rotation_matrices,
+        pose_correctives_regressor=None,
+        template_vertex_positions=template,
+        num_joints=2,
+        num_vertices=3,
+    )
+    self.assertEqual(tuple(result.shape), (2, 3, 3))
+    self.assertEqual(result.dtype, rotation_matrices.dtype)
+
 
 class TestConstantBatch(parameterized.TestCase):
   """Tests for _is_constant_batch."""