diff --git a/mjx/mujoco/mjx/_src/types.py b/mjx/mujoco/mjx/_src/types.py
index 89dab88e..6446cd33 100644
--- a/mjx/mujoco/mjx/_src/types.py
+++ b/mjx/mujoco/mjx/_src/types.py
@@ -933,7 +933,7 @@ class Model(PyTreeNode):
   geom_sameframe: np.ndarray
   geom_dataid: np.ndarray
   geom_group: np.ndarray
-  geom_matid: np.ndarray
+  geom_matid: jax.Array
   geom_priority: np.ndarray
   geom_solmix: jax.Array
   geom_solref: jax.Array