diff --git a/gnm/shape/data/versions/v3_0/gnm_head.npz b/gnm/shape/data/versions/v3_0/gnm_head.npz
index de5f528..d3c45a3 100644
Binary files a/gnm/shape/data/versions/v3_0/gnm_head.npz and b/gnm/shape/data/versions/v3_0/gnm_head.npz differ
diff --git a/gnm/shape/gnm_common_test.py b/gnm/shape/gnm_common_test.py
index 98a914f..4bfb5aa 100644
--- a/gnm/shape/gnm_common_test.py
+++ b/gnm/shape/gnm_common_test.py
@@ -12,13 +12,16 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+"""Tests for the backend-agnostic GNM math."""
+
+from collections.abc import Sequence
+
 from absl.testing import absltest
 from absl.testing import parameterized
 from etils import enp
 from gnm.shape import gnm_common
 import numpy as np
 
-
 # The array backends supported by the GNM math, as (name, xnp module) pairs.
 # Every test is run against all of them so that backend-specific regressions
 # are caught. The xnp modules mirror those used by the per-backend GNM classes
@@ -35,21 +38,21 @@ class ArrayHelpersTest(parameterized.TestCase):
   """Tests for the backend-agnostic array helpers."""
 
   @parameterized.named_parameters(*_XNP_BACKENDS)
-  def test_take_selects_along_axis(self, xnp):
+  def test_take_selects_along_axis(self, xnp: enp.NpModule):
     array = xnp.asarray([[0, 1, 2], [3, 4, 5], [6, 7, 8]], dtype=xnp.float32)
     indices = xnp.asarray([0, 2], dtype=xnp.int32)
     taken = gnm_common.take(array, indices, axis=0, xnp=xnp)
     np.testing.assert_array_equal(np.asarray(taken), [[0, 1, 2], [6, 7, 8]])
 
   @parameterized.named_parameters(*_XNP_BACKENDS)
-  def test_take_selects_along_last_axis(self, xnp):
+  def test_take_selects_along_last_axis(self, xnp: enp.NpModule):
     array = xnp.asarray([[0, 1, 2], [3, 4, 5]], dtype=xnp.float32)
     indices = xnp.asarray([2, 0], dtype=xnp.int32)
     taken = gnm_common.take(array, indices, axis=-1, xnp=xnp)
     np.testing.assert_array_equal(np.asarray(taken), [[2, 0], [5, 3]])
 
   @parameterized.named_parameters(*_XNP_BACKENDS)
-  def test_eye_returns_identity(self, xnp):
+  def test_eye_returns_identity(self, xnp: enp.NpModule):
     reference = xnp.asarray([0.0, 0.0, 0.0, 0.0], dtype=xnp.float32)
     identity = gnm_common.eye(
         3, dtype=xnp.float32, reference_array=reference, xnp=xnp
@@ -58,7 +61,7 @@ def test_eye_returns_identity(self, xnp):
     np.testing.assert_array_equal(np.asarray(identity), np.eye(3))
 
   @parameterized.named_parameters(*_XNP_BACKENDS)
-  def test_reshape_with_batch_dims(self, xnp):
+  def test_reshape_with_batch_dims(self, xnp: enp.NpModule):
     reference = xnp.asarray(np.zeros((2, 5, 3)), dtype=xnp.float32)
     array = xnp.asarray(np.arange(2 * 6), dtype=xnp.float32)
     reshaped = gnm_common.reshape_with_batch_dims(
@@ -70,7 +73,7 @@ def test_reshape_with_batch_dims(self, xnp):
     self.assertEqual(tuple(reshaped.shape), (2, 6))
 
   @parameterized.named_parameters(*_XNP_BACKENDS)
-  def test_zeros_with_batch_dims(self, xnp):
+  def test_zeros_with_batch_dims(self, xnp: enp.NpModule):
     reference = xnp.asarray(np.ones((4, 5, 3)), dtype=xnp.float32)
     zeros = gnm_common.zeros_with_batch_dims(
         reference,
@@ -86,36 +89,42 @@ class AxisAngleToRotationMatrixTest(parameterized.TestCase):
   """Tests for axis_angle_to_rotation_matrix."""
 
   @parameterized.named_parameters(*_XNP_BACKENDS)
-  def test_zero_returns_identity(self, xnp):
+  def test_zero_returns_identity(self, xnp: enp.NpModule):
     axis_angle = xnp.asarray([0.0, 0.0, 0.0], dtype=xnp.float32)
     matrix = np.asarray(gnm_common.axis_angle_to_rotation_matrix(axis_angle))
     np.testing.assert_allclose(matrix, np.eye(3), atol=1e-5)
 
-  @parameterized.named_parameters(*_XNP_BACKENDS)
-  def test_axis_is_invariant(self, xnp):
+  @parameterized.product(
+      xnp=tuple(xnp for _, xnp in _XNP_BACKENDS),
+      case=(
+          ([np.pi / 2, 0.0, 0.0], [1.0, 0.0, 0.0]),
+          ([0.0, np.pi / 2, 0.0], [0.0, 1.0, 0.0]),
+          ([0.0, 0.0, np.pi / 2], [0.0, 0.0, 1.0]),
+      ),
+  )
+  def test_axis_is_invariant(
+      self,
+      xnp: enp.NpModule,
+      case: tuple[Sequence[float], Sequence[float]],
+  ) -> None:
     # Rotating about an axis leaves that axis unchanged.
-    cases = (
-        ([np.pi / 2, 0.0, 0.0], [1.0, 0.0, 0.0]),
-        ([0.0, np.pi / 2, 0.0], [0.0, 1.0, 0.0]),
-        ([0.0, 0.0, np.pi / 2], [0.0, 0.0, 1.0]),
+    axis_angle, axis = case
+    matrix = np.asarray(
+        gnm_common.axis_angle_to_rotation_matrix(
+            xnp.asarray(axis_angle, dtype=xnp.float32)
+        )
     )
-    for axis_angle, axis in cases:
-      matrix = np.asarray(
-          gnm_common.axis_angle_to_rotation_matrix(
-              xnp.asarray(axis_angle, dtype=xnp.float32)
-          )
-      )
-      np.testing.assert_allclose(matrix @ np.asarray(axis), axis, atol=1e-5)
+    np.testing.assert_allclose(matrix @ np.asarray(axis), axis, atol=1e-5)
 
   @parameterized.named_parameters(*_XNP_BACKENDS)
-  def test_ninety_degrees_about_z(self, xnp):
+  def test_ninety_degrees_about_z(self, xnp: enp.NpModule):
     axis_angle = xnp.asarray([0.0, 0.0, np.pi / 2], dtype=xnp.float32)
     matrix = np.asarray(gnm_common.axis_angle_to_rotation_matrix(axis_angle))
     expected = np.array([[0.0, -1.0, 0.0], [1.0, 0.0, 0.0], [0.0, 0.0, 1.0]])
     np.testing.assert_allclose(matrix, expected, atol=1e-5)
 
   @parameterized.named_parameters(*_XNP_BACKENDS)
-  def test_matrix_is_orthonormal_with_unit_determinant(self, xnp):
+  def test_matrix_is_orthonormal_with_unit_determinant(self, xnp: enp.NpModule):
     rng = np.random.default_rng(0)
     axis_angle = rng.uniform(-3.0, 3.0, size=(8, 3))
     matrices = np.asarray(
@@ -131,7 +140,7 @@ def test_matrix_is_orthonormal_with_unit_determinant(self, xnp):
     np.testing.assert_allclose(np.linalg.det(matrices), np.ones(8), atol=1e-4)
 
   @parameterized.named_parameters(*_XNP_BACKENDS)
-  def test_preserves_batch_shape(self, xnp):
+  def test_preserves_batch_shape(self, xnp: enp.NpModule):
     axis_angle = xnp.asarray(np.zeros((2, 4, 3)), dtype=xnp.float32)
     matrices = gnm_common.axis_angle_to_rotation_matrix(axis_angle)
     self.assertEqual(tuple(matrices.shape), (2, 4, 3, 3))
@@ -147,7 +156,7 @@ def setUp(self):
     self.parents = [0, 0]
 
   @parameterized.named_parameters(*_XNP_BACKENDS)
-  def test_rest_pose_places_joints_at_bind_positions(self, xnp):
+  def test_rest_pose_places_joints_at_bind_positions(self, xnp: enp.NpModule):
     transforms = np.asarray(
         gnm_common.joint_transforms_world(
             xnp.asarray(self.joints, dtype=xnp.float32),
@@ -166,7 +175,7 @@ def test_rest_pose_places_joints_at_bind_positions(self, xnp):
     )
 
   @parameterized.named_parameters(*_XNP_BACKENDS)
-  def test_global_translation_shifts_all_joints(self, xnp):
+  def test_global_translation_shifts_all_joints(self, xnp: enp.NpModule):
     transforms = np.asarray(
         gnm_common.joint_transforms_world(
             xnp.asarray(self.joints, dtype=xnp.float32),
@@ -182,7 +191,7 @@ def test_global_translation_shifts_all_joints(self, xnp):
     )
 
   @parameterized.named_parameters(*_XNP_BACKENDS)
-  def test_root_rotation_propagates_to_child(self, xnp):
+  def test_root_rotation_propagates_to_child(self, xnp: enp.NpModule):
     # Rotate the root 90 degrees about z; the child at (1, 0, 0) maps to
     # (0, 1, 0).
     rotations = xnp.asarray(
@@ -213,7 +222,7 @@ def setUp(self):
     self.vertices = [[0.5, 0.0, 0.0], [0.0, 0.5, 0.0], [1.0, 1.0, 0.0]]
 
   @parameterized.named_parameters(*_XNP_BACKENDS)
-  def test_rest_pose_is_identity(self, xnp):
+  def test_rest_pose_is_identity(self, xnp: enp.NpModule):
     posed = np.asarray(
         gnm_common.linear_blend_skinning(
             xnp.asarray(self.vertices, dtype=xnp.float32),
@@ -228,7 +237,7 @@ def test_rest_pose_is_identity(self, xnp):
     np.testing.assert_allclose(posed, self.vertices, atol=1e-5)
 
   @parameterized.named_parameters(*_XNP_BACKENDS)
-  def test_global_translation_translates_vertices(self, xnp):
+  def test_global_translation_translates_vertices(self, xnp: enp.NpModule):
     translation = [10.0, 0.0, 0.0]
     posed = np.asarray(
         gnm_common.linear_blend_skinning(
@@ -245,7 +254,7 @@ def test_global_translation_translates_vertices(self, xnp):
     )
 
   @parameterized.named_parameters(*_XNP_BACKENDS)
-  def test_preserves_batch_shape(self, xnp):
+  def test_preserves_batch_shape(self, xnp: enp.NpModule):
     batch_vertices = np.array(np.broadcast_to(self.vertices, (4, 3, 3)))
     batch_joints = np.array(np.broadcast_to(self.joints, (4, 2, 3)))
     posed = gnm_common.linear_blend_skinning(
@@ -263,7 +272,9 @@ class BindPoseTest(parameterized.TestCase):
   """Tests for the bind-pose vertex and joint helpers."""
 
   @parameterized.named_parameters(*_XNP_BACKENDS)
-  def test_vertex_positions_none_params_return_template(self, xnp):
+  def test_vertex_positions_none_params_return_template(
+      self, xnp: enp.NpModule
+  ):
     template = xnp.asarray(np.ones((3, 3)), dtype=xnp.float32)
     identity_basis = xnp.asarray(np.zeros((2, 3, 3)), dtype=xnp.float32)
     expression_basis = xnp.asarray(np.zeros((2, 3, 3)), dtype=xnp.float32)
@@ -273,7 +284,9 @@ def test_vertex_positions_none_params_return_template(self, xnp):
     np.testing.assert_allclose(np.asarray(result), np.ones((3, 3)), atol=1e-5)
 
   @parameterized.named_parameters(*_XNP_BACKENDS)
-  def test_vertex_positions_apply_identity_and_expression(self, xnp):
+  def test_vertex_positions_apply_identity_and_expression(
+      self, xnp: enp.NpModule
+  ):
     template = xnp.asarray(np.zeros((2, 3)), dtype=xnp.float32)
     # (I=1, V=2, 3)
     identity_basis = xnp.asarray(
@@ -294,7 +307,9 @@ def test_vertex_positions_apply_identity_and_expression(self, xnp):
     np.testing.assert_allclose(np.asarray(result), expected, atol=1e-5)
 
   @parameterized.named_parameters(*_XNP_BACKENDS)
-  def test_joint_positions_none_identity_returns_template(self, xnp):
+  def test_joint_positions_none_identity_returns_template(
+      self, xnp: enp.NpModule
+  ):
     template = xnp.asarray(
         [[0.0, 0.0, 0.0], [1.0, 0.0, 0.0]], dtype=xnp.float32
     )
@@ -305,7 +320,7 @@ def test_joint_positions_none_identity_returns_template(self, xnp):
     )
 
   @parameterized.named_parameters(*_XNP_BACKENDS)
-  def test_joint_positions_apply_identity(self, xnp):
+  def test_joint_positions_apply_identity(self, xnp: enp.NpModule):
     template = xnp.asarray(np.zeros((2, 3)), dtype=xnp.float32)
     # (I=1, J=2, 3)
     joint_basis = xnp.asarray(
@@ -323,7 +338,7 @@ class PoseCorrectivesTest(parameterized.TestCase):
   """Tests for compute_pose_correctives."""
 
   @parameterized.named_parameters(*_XNP_BACKENDS)
-  def test_none_inputs_return_zeros(self, xnp):
+  def test_none_inputs_return_zeros(self, xnp: enp.NpModule):
     template = xnp.asarray(np.ones((3, 3)), dtype=xnp.float32)
     result = gnm_common.compute_pose_correctives(
         None, None, template, num_joints=2, num_vertices=3
@@ -332,7 +347,7 @@ def test_none_inputs_return_zeros(self, xnp):
     np.testing.assert_array_equal(np.asarray(result), np.zeros((3, 3)))
 
   @parameterized.named_parameters(*_XNP_BACKENDS)
-  def test_zero_rotations_produce_no_correctives(self, xnp):
+  def test_zero_rotations_produce_no_correctives(self, xnp: enp.NpModule):
     template = xnp.asarray(np.ones((3, 3)), dtype=xnp.float32)
     regressor = xnp.asarray(np.ones((2 * 9, 3 * 3)), dtype=xnp.float32)
     result = gnm_common.compute_pose_correctives(
@@ -346,7 +361,7 @@ def test_zero_rotations_produce_no_correctives(self, xnp):
     np.testing.assert_allclose(np.asarray(result), np.zeros((3, 3)), atol=1e-5)
 
   @parameterized.named_parameters(*_XNP_BACKENDS)
-  def test_preserves_batch_shape(self, xnp):
+  def test_preserves_batch_shape(self, xnp: enp.NpModule):
     template = xnp.asarray(np.ones((3, 3)), dtype=xnp.float32)
     regressor = xnp.asarray(np.ones((2 * 9, 3 * 3)), dtype=xnp.float32)
     result = gnm_common.compute_pose_correctives(
diff --git a/gnm/shape/gnm_landmarks_test.py b/gnm/shape/gnm_landmarks_test.py
index 37b3fa0..2cbe4b0 100644
--- a/gnm/shape/gnm_landmarks_test.py
+++ b/gnm/shape/gnm_landmarks_test.py
@@ -15,6 +15,7 @@
 """Tests for validating and loading GNM landmarks configurations."""
 
 from absl.testing import absltest
+
 from gnm.shape import gnm_landmarks
 from gnm.shape.data.versions import gnm_specs
 import numpy as np