Import google-deepmind/mujoco_warp from GitHub.

PiperOrigin-RevId: 899731493
Change-Id: Ib563023047e85db1aff56e78724dc28d769ef91e

mjx/mujoco/mjx/_src/io.py

@@ -456,6 +456,9 @@ def _put_model_warp(
     if not hasattr(mw, k) or k in ('stat', 'opt'):
       continue
     field = _wp_to_np_type(getattr(mw, k), k)
+    if k == 'geom_dataid' and field.ndim > 1:
+      # Batched geom_dataid is not supported in MJX.
+      field = field[0]
     fields[k] = field
 
   impl_fields = {}

mjx/mujoco/mjx/third_party/mujoco_warp/_src/benchmark.py

@@ -41,18 +41,18 @@ def _sum(stack1, stack2):
 @wp.kernel
 def ctrl_noise(
   # Model:
-  opt_timestep: wp.array(dtype=float),
-  actuator_ctrllimited: wp.array(dtype=bool),
-  actuator_ctrlrange: wp.array2d(dtype=wp.vec2),
+  opt_timestep: wp.array[float],
+  actuator_ctrllimited: wp.array[bool],
+  actuator_ctrlrange: wp.array2d[wp.vec2],
   # Data in:
-  ctrl_in: wp.array2d(dtype=float),
+  ctrl_in: wp.array2d[float],
   # In:
-  ctrl_center: wp.array1d(dtype=float),
+  ctrl_center: wp.array[float],
   step: int,
   ctrlnoisestd: float,
   ctrlnoiserate: float,
   # Data out:
-  ctrl_out: wp.array2d(dtype=float),
+  ctrl_out: wp.array2d[float],
 ):
   worldid, actid = wp.tid()
 

mjx/mujoco/mjx/third_party/mujoco_warp/_src/block_cholesky.py

@@ -23,10 +23,10 @@ def create_blocked_cholesky_func(block_size: int):
   @wp.func
   def blocked_cholesky_func(
     # In:
-    A: wp.array(dtype=float, ndim=2),
+    A: wp.array2d[float],
     matrix_size: int,
     # Out:
-    L: wp.array(dtype=float, ndim=2),
+    L: wp.array2d[float],
   ):
     """Computes the Cholesky factorization of a symmetric positive definite matrix A in blocks.
 
@@ -68,11 +68,11 @@ def create_blocked_cholesky_solve_func(block_size: int, matrix_size_static: int)
   @wp.func
   def blocked_cholesky_solve_func(
     # In:
-    L: wp.array(dtype=float, ndim=2),
-    b: wp.array(dtype=float, ndim=2),
+    L: wp.array2d[float],
+    b: wp.array2d[float],
     matrix_size: int,
     # Out:
-    x: wp.array(dtype=float, ndim=2),
+    x: wp.array2d[float],
   ):
     """Block Cholesky factorization and solve.
 

mjx/mujoco/mjx/third_party/mujoco_warp/_src/bvh.py

@@ -173,21 +173,21 @@ def _compute_cylinder_bounds(
 @wp.kernel
 def _compute_bvh_bounds(
   # Model:
-  geom_type: wp.array(dtype=int),
-  geom_dataid: wp.array(dtype=int),
-  geom_size: wp.array2d(dtype=wp.vec3),
+  geom_type: wp.array[int],
+  geom_dataid: wp.array2d[int],
+  geom_size: wp.array2d[wp.vec3],
   # Data in:
-  geom_xpos_in: wp.array2d(dtype=wp.vec3),
-  geom_xmat_in: wp.array2d(dtype=wp.mat33),
+  geom_xpos_in: wp.array2d[wp.vec3],
+  geom_xmat_in: wp.array2d[wp.mat33],
   # In:
   bvh_ngeom: int,
-  enabled_geom_ids: wp.array(dtype=int),
-  mesh_bounds_size: wp.array(dtype=wp.vec3),
-  hfield_bounds_size: wp.array(dtype=wp.vec3),
+  enabled_geom_ids: wp.array[int],
+  mesh_bounds_size: wp.array[wp.vec3],
+  hfield_bounds_size: wp.array[wp.vec3],
   # Out:
-  lower_out: wp.array(dtype=wp.vec3),
-  upper_out: wp.array(dtype=wp.vec3),
-  group_out: wp.array(dtype=int),
+  lower_out: wp.array[wp.vec3],
+  upper_out: wp.array[wp.vec3],
+  group_out: wp.array[int],
 ):
   worldid, geom_local_id = wp.tid()
   geom_id = enabled_geom_ids[geom_local_id]
@@ -205,8 +205,13 @@ def _compute_bvh_bounds(
   elif type == GeomType.PLANE:
     lower_bound, upper_bound = _compute_plane_bounds(pos, rot, size)
   elif type == GeomType.MESH:
-    size = mesh_bounds_size[geom_dataid[geom_id]]
-    lower_bound, upper_bound = _compute_box_bounds(pos, rot, size)
+    did = geom_dataid[worldid % geom_dataid.shape[0], geom_id]
+    if did >= 0:
+      size = mesh_bounds_size[did]
+      lower_bound, upper_bound = _compute_box_bounds(pos, rot, size)
+    else:
+      lower_bound = pos
+      upper_bound = pos
   elif type == GeomType.ELLIPSOID:
     lower_bound, upper_bound = _compute_ellipsoid_bounds(pos, rot, size)
   elif type == GeomType.CYLINDER:
@@ -214,9 +219,14 @@ def _compute_bvh_bounds(
   elif type == GeomType.BOX:
     lower_bound, upper_bound = _compute_box_bounds(pos, rot, size)
   elif type == GeomType.HFIELD:
-    size = hfield_bounds_size[geom_dataid[geom_id]]
-    hfield_center = pos + rot[:, 2] * size[2]
-    lower_bound, upper_bound = _compute_box_bounds(hfield_center, rot, size)
+    did = geom_dataid[worldid % geom_dataid.shape[0], geom_id]
+    if did >= 0:
+      size = hfield_bounds_size[did]
+      hfield_center = pos + rot[:, 2] * size[2]
+      lower_bound, upper_bound = _compute_box_bounds(hfield_center, rot, size)
+    else:
+      lower_bound = pos
+      upper_bound = pos
 
   lower_out[worldid * bvh_ngeom + geom_local_id] = lower_bound
   upper_out[worldid * bvh_ngeom + geom_local_id] = upper_bound
@@ -228,7 +238,7 @@ def compute_bvh_group_roots(
   # In:
   bvh_id: wp.uint64,
   # Out:
-  group_root_out: wp.array(dtype=int),
+  group_root_out: wp.array[int],
 ):
   tid = wp.tid()
   root = wp.bvh_get_group_root(bvh_id, tid)
@@ -238,21 +248,21 @@ def compute_bvh_group_roots(
 @wp.kernel
 def _compute_flex_bvh_bounds(
   # Model:
-  flex_vertadr: wp.array(dtype=int),
-  flex_vertnum: wp.array(dtype=int),
-  flex_edge: wp.array(dtype=wp.vec2i),
-  flex_radius: wp.array(dtype=float),
+  flex_vertadr: wp.array[int],
+  flex_vertnum: wp.array[int],
+  flex_edge: wp.array[wp.vec2i],
+  flex_radius: wp.array[float],
   # Data in:
-  flexvert_xpos_in: wp.array2d(dtype=wp.vec3),
+  flexvert_xpos_in: wp.array2d[wp.vec3],
   # In:
-  flex_geom_flexid: wp.array(dtype=int),
-  flex_geom_edgeid: wp.array(dtype=int),
+  flex_geom_flexid: wp.array[int],
+  flex_geom_edgeid: wp.array[int],
   bvh_ngeom: int,
   total_bvh_size: int,
   # Out:
-  lower_out: wp.array(dtype=wp.vec3),
-  upper_out: wp.array(dtype=wp.vec3),
-  group_out: wp.array(dtype=int),
+  lower_out: wp.array[wp.vec3],
+  upper_out: wp.array[wp.vec3],
+  group_out: wp.array[int],
 ):
   worldid, flexlocalid = wp.tid()
 
@@ -289,7 +299,7 @@ def build_scene_bvh(mjm: mujoco.MjModel, mjd: mujoco.MjData, rc: RenderContext,
   total_bvh_size = rc.bvh_ngeom + rc.bvh_nflexgeom
 
   geom_type = wp.array(mjm.geom_type, dtype=int)
-  geom_dataid = wp.array(mjm.geom_dataid, dtype=int)
+  geom_dataid = wp.array(np.tile(mjm.geom_dataid, (nworld, 1)), dtype=int)
   geom_size = wp.array(np.tile(mjm.geom_size[np.newaxis, :, :], (nworld, 1, 1)), dtype=wp.vec3)
   geom_xpos = wp.array(np.tile(mjd.geom_xpos[np.newaxis, :, :], (nworld, 1, 1)), dtype=wp.vec3)
   geom_xmat = wp.array(np.tile(mjd.geom_xmat.reshape(mjm.ngeom, 3, 3)[np.newaxis, :, :, :], (nworld, 1, 1, 1)), dtype=wp.mat33)
@@ -599,16 +609,16 @@ def build_hfield_bvh(
 def accumulate_flex_vertex_normals(
   # Model:
   nflex: int,
-  flex_dim: wp.array(dtype=int),
-  flex_vertadr: wp.array(dtype=int),
-  flex_elemadr: wp.array(dtype=int),
-  flex_elemnum: wp.array(dtype=int),
-  flex_elemdataadr: wp.array(dtype=int),
-  flex_elem: wp.array(dtype=int),
+  flex_dim: wp.array[int],
+  flex_vertadr: wp.array[int],
+  flex_elemadr: wp.array[int],
+  flex_elemnum: wp.array[int],
+  flex_elemdataadr: wp.array[int],
+  flex_elem: wp.array[int],
   # Data in:
-  flexvert_xpos_in: wp.array2d(dtype=wp.vec3),
+  flexvert_xpos_in: wp.array2d[wp.vec3],
   # Out:
-  flexvert_norm_out: wp.array2d(dtype=wp.vec3),
+  flexvert_norm_out: wp.array2d[wp.vec3],
 ):
   """Accumulate per-vertex normals by summing adjacent face normals."""
   worldid, elemid = wp.tid()
@@ -644,7 +654,7 @@ def accumulate_flex_vertex_normals(
 @wp.kernel
 def normalize_vertex_normals(
   # Out:
-  flexvert_norm_out: wp.array2d(dtype=wp.vec3),
+  flexvert_norm_out: wp.array2d[wp.vec3],
 ):
   """Normalize accumulated vertex normals."""
   worldid, vertid = wp.tid()
@@ -654,20 +664,20 @@ def normalize_vertex_normals(
 @wp.kernel
 def _build_flex_2d_elements(
   # Model:
-  flex_elem: wp.array(dtype=int),
+  flex_elem: wp.array[int],
   # Data in:
-  flexvert_xpos_in: wp.array2d(dtype=wp.vec3),
+  flexvert_xpos_in: wp.array2d[wp.vec3],
   # In:
-  flexvert_norm_in: wp.array2d(dtype=wp.vec3),
+  flexvert_norm_in: wp.array2d[wp.vec3],
   elem_adr: int,
   vert_adr: int,
   face_offset: int,
   radius: float,
   nfaces: int,
   # Out:
-  face_point_out: wp.array(dtype=wp.vec3),
-  face_index_out: wp.array(dtype=int),
-  group_out: wp.array(dtype=int),
+  face_point_out: wp.array[wp.vec3],
+  face_index_out: wp.array[int],
+  group_out: wp.array[int],
 ):
   """Create faces from 2D flex elements (triangles).
 
@@ -728,20 +738,20 @@ def _build_flex_2d_elements(
 @wp.kernel
 def _build_flex_2d_sides(
   # Model:
-  flex_shell: wp.array(dtype=int),
+  flex_shell: wp.array[int],
   # Data in:
-  flexvert_xpos_in: wp.array2d(dtype=wp.vec3),
+  flexvert_xpos_in: wp.array2d[wp.vec3],
   # In:
-  flexvert_norm_in: wp.array2d(dtype=wp.vec3),
+  flexvert_norm_in: wp.array2d[wp.vec3],
   shell_adr: int,
   vert_adr: int,
   face_offset: int,
   radius: float,
   nface: int,
   # Out:
-  face_point_out: wp.array(dtype=wp.vec3),
-  face_index_out: wp.array(dtype=int),
-  group_out: wp.array(dtype=int),
+  face_point_out: wp.array[wp.vec3],
+  face_index_out: wp.array[int],
+  group_out: wp.array[int],
 ):
   """Create side faces from 2D flex shell fragments.
 
@@ -790,18 +800,18 @@ def _build_flex_2d_sides(
 @wp.kernel
 def _build_flex_3d_shells(
   # Model:
-  flex_shell: wp.array(dtype=int),
+  flex_shell: wp.array[int],
   # Data in:
-  flexvert_xpos_in: wp.array2d(dtype=wp.vec3),
+  flexvert_xpos_in: wp.array2d[wp.vec3],
   # In:
   shell_adr: int,
   vert_adr: int,
   face_offset: int,
   nface: int,
   # Out:
-  face_point_out: wp.array(dtype=wp.vec3),
-  face_index_out: wp.array(dtype=int),
-  group_out: wp.array(dtype=int),
+  face_point_out: wp.array[wp.vec3],
+  face_index_out: wp.array[int],
+  group_out: wp.array[int],
 ):
   """Create faces from 3D flex shell fragments (triangles).
 
@@ -836,22 +846,22 @@ def _build_flex_3d_shells(
 @wp.kernel
 def _update_flex_2d_face_points(
   # Model:
-  flex_vertadr: wp.array(dtype=int),
-  flex_elemnum: wp.array(dtype=int),
-  flex_elemdataadr: wp.array(dtype=int),
-  flex_shelldataadr: wp.array(dtype=int),
-  flex_elem: wp.array(dtype=int),
-  flex_shell: wp.array(dtype=int),
-  flex_radius: wp.array(dtype=float),
+  flex_vertadr: wp.array[int],
+  flex_elemnum: wp.array[int],
+  flex_elemdataadr: wp.array[int],
+  flex_shelldataadr: wp.array[int],
+  flex_elem: wp.array[int],
+  flex_shell: wp.array[int],
+  flex_radius: wp.array[float],
   # Data in:
-  flexvert_xpos_in: wp.array2d(dtype=wp.vec3),
+  flexvert_xpos_in: wp.array2d[wp.vec3],
   # In:
-  flexvert_norm_in: wp.array2d(dtype=wp.vec3),
+  flexvert_norm_in: wp.array2d[wp.vec3],
   flex_id: int,
   nface: int,
   smooth: bool,
   # Out:
-  face_point_out: wp.array(dtype=wp.vec3),
+  face_point_out: wp.array[wp.vec3],
 ):
   worldid, workid = wp.tid()
 
@@ -935,16 +945,16 @@ def _update_flex_2d_face_points(
 @wp.kernel
 def _update_flex_3d_face_points(
   # Model:
-  flex_vertadr: wp.array(dtype=int),
-  flex_shelldataadr: wp.array(dtype=int),
-  flex_shell: wp.array(dtype=int),
+  flex_vertadr: wp.array[int],
+  flex_shelldataadr: wp.array[int],
+  flex_shell: wp.array[int],
   # Data in:
-  flexvert_xpos_in: wp.array2d(dtype=wp.vec3),
+  flexvert_xpos_in: wp.array2d[wp.vec3],
   # In:
   flex_id: int,
   nface: int,
   # Out:
-  face_point_out: wp.array(dtype=wp.vec3),
+  face_point_out: wp.array[wp.vec3],
 ):
   worldid, shellid = wp.tid()
 

mjx/mujoco/mjx/third_party/mujoco_warp/_src/collision_convex.py

@@ -55,20 +55,20 @@ mat_maxconpair = wp.types.matrix(shape=(MJ_MAXCONPAIR, 3), dtype=float)
 @wp.func
 def _hfield_filter(
   # Model:
-  geom_type: wp.array(dtype=int),
-  geom_dataid: wp.array(dtype=int),
-  geom_size: wp.array2d(dtype=wp.vec3),
-  geom_rbound: wp.array2d(dtype=float),
-  geom_margin: wp.array2d(dtype=float),
-  mesh_vertadr: wp.array(dtype=int),
-  mesh_vertnum: wp.array(dtype=int),
-  mesh_graphadr: wp.array(dtype=int),
-  mesh_vert: wp.array(dtype=wp.vec3),
-  mesh_graph: wp.array(dtype=int),
-  hfield_size: wp.array(dtype=wp.vec4),
+  geom_type: wp.array[int],
+  geom_dataid: wp.array2d[int],
+  geom_size: wp.array2d[wp.vec3],
+  geom_rbound: wp.array2d[float],
+  geom_margin: wp.array2d[float],
+  mesh_vertadr: wp.array[int],
+  mesh_vertnum: wp.array[int],
+  mesh_graphadr: wp.array[int],
+  mesh_vert: wp.array[wp.vec3],
+  mesh_graph: wp.array[int],
+  hfield_size: wp.array[wp.vec4],
   # Data in:
-  geom_xpos_in: wp.array2d(dtype=wp.vec3),
-  geom_xmat_in: wp.array2d(dtype=wp.mat33),
+  geom_xpos_in: wp.array2d[wp.vec3],
+  geom_xmat_in: wp.array2d[wp.mat33],
   # In:
   worldid: int,
   g1: int,
@@ -79,7 +79,8 @@ def _hfield_filter(
   See MuJoCo mjc_ConvexHField.
   """
   # height field info
-  hfdataid = geom_dataid[g1]
+  dataid_setid = worldid % geom_dataid.shape[0]
+  hfdataid = geom_dataid[dataid_setid, g1]
   size1 = hfield_size[hfdataid]
 
   # geom info
@@ -124,7 +125,7 @@ def _hfield_filter(
 
   # load mesh vertex data for support function queries
   if geomtype2 == GeomType.MESH:
-    dataid = geom_dataid[g2]
+    dataid = geom_dataid[dataid_setid, g2]
     geom2.vertadr = wp.where(dataid >= 0, mesh_vertadr[dataid], -1)
     geom2.vertnum = wp.where(dataid >= 0, mesh_vertnum[dataid], -1)
     geom2.graphadr = wp.where(dataid >= 0, mesh_graphadr[dataid], -1)
@@ -167,78 +168,78 @@ def ccd_hfield_kernel_builder(
   @wp.kernel(module="unique", enable_backward=False)
   def ccd_hfield_kernel(
     # Model:
-    opt_ccd_tolerance: wp.array(dtype=float),
-    geom_type: wp.array(dtype=int),
-    geom_condim: wp.array(dtype=int),
-    geom_dataid: wp.array(dtype=int),
-    geom_priority: wp.array(dtype=int),
-    geom_solmix: wp.array2d(dtype=float),
-    geom_solref: wp.array2d(dtype=wp.vec2),
-    geom_solimp: wp.array2d(dtype=vec5),
-    geom_size: wp.array2d(dtype=wp.vec3),
-    geom_rbound: wp.array2d(dtype=float),
-    geom_friction: wp.array2d(dtype=wp.vec3),
-    geom_margin: wp.array2d(dtype=float),
-    geom_gap: wp.array2d(dtype=float),
-    mesh_vertadr: wp.array(dtype=int),
-    mesh_vertnum: wp.array(dtype=int),
-    mesh_graphadr: wp.array(dtype=int),
-    mesh_vert: wp.array(dtype=wp.vec3),
-    mesh_graph: wp.array(dtype=int),
-    mesh_polynum: wp.array(dtype=int),
-    mesh_polyadr: wp.array(dtype=int),
-    mesh_polynormal: wp.array(dtype=wp.vec3),
-    mesh_polyvertadr: wp.array(dtype=int),
-    mesh_polyvertnum: wp.array(dtype=int),
-    mesh_polyvert: wp.array(dtype=int),
-    mesh_polymapadr: wp.array(dtype=int),
-    mesh_polymapnum: wp.array(dtype=int),
-    mesh_polymap: wp.array(dtype=int),
-    hfield_size: wp.array(dtype=wp.vec4),
-    hfield_nrow: wp.array(dtype=int),
-    hfield_ncol: wp.array(dtype=int),
-    hfield_adr: wp.array(dtype=int),
-    hfield_data: wp.array(dtype=float),
-    pair_dim: wp.array(dtype=int),
-    pair_solref: wp.array2d(dtype=wp.vec2),
-    pair_solreffriction: wp.array2d(dtype=wp.vec2),
-    pair_solimp: wp.array2d(dtype=vec5),
-    pair_margin: wp.array2d(dtype=float),
-    pair_gap: wp.array2d(dtype=float),
-    pair_friction: wp.array2d(dtype=vec5),
+    opt_ccd_tolerance: wp.array[float],
+    geom_type: wp.array[int],
+    geom_condim: wp.array[int],
+    geom_dataid: wp.array2d[int],
+    geom_priority: wp.array[int],
+    geom_solmix: wp.array2d[float],
+    geom_solref: wp.array2d[wp.vec2],
+    geom_solimp: wp.array2d[vec5],
+    geom_size: wp.array2d[wp.vec3],
+    geom_rbound: wp.array2d[float],
+    geom_friction: wp.array2d[wp.vec3],
+    geom_margin: wp.array2d[float],
+    geom_gap: wp.array2d[float],
+    mesh_vertadr: wp.array[int],
+    mesh_vertnum: wp.array[int],
+    mesh_graphadr: wp.array[int],
+    mesh_vert: wp.array[wp.vec3],
+    mesh_graph: wp.array[int],
+    mesh_polynum: wp.array[int],
+    mesh_polyadr: wp.array[int],
+    mesh_polynormal: wp.array[wp.vec3],
+    mesh_polyvertadr: wp.array[int],
+    mesh_polyvertnum: wp.array[int],
+    mesh_polyvert: wp.array[int],
+    mesh_polymapadr: wp.array[int],
+    mesh_polymapnum: wp.array[int],
+    mesh_polymap: wp.array[int],
+    hfield_size: wp.array[wp.vec4],
+    hfield_nrow: wp.array[int],
+    hfield_ncol: wp.array[int],
+    hfield_adr: wp.array[int],
+    hfield_data: wp.array[float],
+    pair_dim: wp.array[int],
+    pair_solref: wp.array2d[wp.vec2],
+    pair_solreffriction: wp.array2d[wp.vec2],
+    pair_solimp: wp.array2d[vec5],
+    pair_margin: wp.array2d[float],
+    pair_gap: wp.array2d[float],
+    pair_friction: wp.array2d[vec5],
     # Data in:
-    geom_xpos_in: wp.array2d(dtype=wp.vec3),
-    geom_xmat_in: wp.array2d(dtype=wp.mat33),
+    geom_xpos_in: wp.array2d[wp.vec3],
+    geom_xmat_in: wp.array2d[wp.mat33],
     naconmax_in: int,
     naccdmax_in: int,
-    ncollision_in: wp.array(dtype=int),
+    ncollision_in: wp.array[int],
     # In:
-    collision_pair_in: wp.array(dtype=wp.vec2i),
-    collision_pairid_in: wp.array(dtype=wp.vec2i),
-    collision_worldid_in: wp.array(dtype=int),
-    epa_vert_in: wp.array2d(dtype=wp.vec3),
-    epa_vert_index_in: wp.array2d(dtype=int),
-    epa_face_in: wp.array2d(dtype=int),
-    epa_pr_in: wp.array2d(dtype=wp.vec3),
-    epa_norm2_in: wp.array2d(dtype=float),
-    epa_horizon_in: wp.array2d(dtype=int),
-    nccd_in: wp.array(dtype=int),
+    collision_pair_in: wp.array[wp.vec2i],
+    collision_pairid_in: wp.array[wp.vec2i],
+    collision_worldid_in: wp.array[int],
+    epa_vert_in: wp.array2d[wp.vec3],
+    epa_vert_index_in: wp.array2d[int],
+    epa_face_in: wp.array2d[int],
+    epa_pr_in: wp.array2d[wp.vec3],
+    epa_norm2_in: wp.array2d[float],
+    epa_horizon_in: wp.array2d[int],
+    nccd_in: wp.array[int],
     # Data out:
-    contact_dist_out: wp.array(dtype=float),
-    contact_pos_out: wp.array(dtype=wp.vec3),
-    contact_frame_out: wp.array(dtype=wp.mat33),
-    contact_includemargin_out: wp.array(dtype=float),
-    contact_friction_out: wp.array(dtype=vec5),
-    contact_solref_out: wp.array(dtype=wp.vec2),
-    contact_solreffriction_out: wp.array(dtype=wp.vec2),
-    contact_solimp_out: wp.array(dtype=vec5),
-    contact_dim_out: wp.array(dtype=int),
-    contact_geom_out: wp.array(dtype=wp.vec2i),
-    contact_efc_address_out: wp.array2d(dtype=int),
-    contact_worldid_out: wp.array(dtype=int),
-    contact_type_out: wp.array(dtype=int),
-    contact_geomcollisionid_out: wp.array(dtype=int),
-    nacon_out: wp.array(dtype=int),
+    contact_dist_out: wp.array[float],
+    contact_pos_out: wp.array[wp.vec3],
+    contact_frame_out: wp.array[wp.mat33],
+    contact_includemargin_out: wp.array[float],
+    contact_friction_out: wp.array[vec5],
+    contact_solref_out: wp.array[wp.vec2],
+    contact_solreffriction_out: wp.array[wp.vec2],
+    contact_solimp_out: wp.array[vec5],
+    contact_dim_out: wp.array[int],
+    contact_geom_out: wp.array[wp.vec2i],
+    contact_efc_address_out: wp.array2d[int],
+    contact_worldid_out: wp.array[int],
+    contact_type_out: wp.array[int],
+    contact_geomcollisionid_out: wp.array[int],
+    nacon_out: wp.array[int],
   ):
     collisionid = wp.tid()
     if collisionid >= ncollision_in[0]:
@@ -339,7 +340,7 @@ def ccd_hfield_kernel_builder(
     geom1.rot = wp.identity(n=3, dtype=float)
 
     # see MuJoCo mjc_ConvexHField
-    geom1_dataid = geom_dataid[g1]
+    geom1_dataid = geom_dataid[worldid % geom_dataid.shape[0], g1]
 
     # height field subgrid
     nrow = hfield_nrow[geom1_dataid]
@@ -710,27 +711,27 @@ def ccd_kernel_builder(
   @wp.func
   def eval_ccd_write_contact(
     # Model:
-    opt_ccd_tolerance: wp.array(dtype=float),
+    opt_ccd_tolerance: wp.array[float],
     # Data in:
     naconmax_in: int,
     # In:
-    epa_vert_in: wp.array2d(dtype=wp.vec3),
-    epa_vert_index_in: wp.array2d(dtype=int),
-    epa_face_in: wp.array2d(dtype=int),
-    epa_pr_in: wp.array2d(dtype=wp.vec3),
-    epa_norm2_in: wp.array2d(dtype=float),
-    epa_horizon_in: wp.array2d(dtype=int),
-    multiccd_polygon_in: wp.array2d(dtype=wp.vec3),
-    multiccd_clipped_in: wp.array2d(dtype=wp.vec3),
-    multiccd_pnormal_in: wp.array2d(dtype=wp.vec3),
-    multiccd_pdist_in: wp.array2d(dtype=float),
-    multiccd_idx1_in: wp.array2d(dtype=int),
-    multiccd_idx2_in: wp.array2d(dtype=int),
-    multiccd_n1_in: wp.array2d(dtype=wp.vec3),
-    multiccd_n2_in: wp.array2d(dtype=wp.vec3),
-    multiccd_endvert_in: wp.array2d(dtype=wp.vec3),
-    multiccd_face1_in: wp.array2d(dtype=wp.vec3),
-    multiccd_face2_in: wp.array2d(dtype=wp.vec3),
+    epa_vert_in: wp.array2d[wp.vec3],
+    epa_vert_index_in: wp.array2d[int],
+    epa_face_in: wp.array2d[int],
+    epa_pr_in: wp.array2d[wp.vec3],
+    epa_norm2_in: wp.array2d[float],
+    epa_horizon_in: wp.array2d[int],
+    multiccd_polygon_in: wp.array2d[wp.vec3],
+    multiccd_clipped_in: wp.array2d[wp.vec3],
+    multiccd_pnormal_in: wp.array2d[wp.vec3],
+    multiccd_pdist_in: wp.array2d[float],
+    multiccd_idx1_in: wp.array2d[int],
+    multiccd_idx2_in: wp.array2d[int],
+    multiccd_n1_in: wp.array2d[wp.vec3],
+    multiccd_n2_in: wp.array2d[wp.vec3],
+    multiccd_endvert_in: wp.array2d[wp.vec3],
+    multiccd_face1_in: wp.array2d[wp.vec3],
+    multiccd_face2_in: wp.array2d[wp.vec3],
     geom1: Geom,
     geom2: Geom,
     geoms: wp.vec2i,
@@ -747,21 +748,21 @@ def ccd_kernel_builder(
     x2: wp.vec3,
     pairid: wp.vec2i,
     # Data out:
-    contact_dist_out: wp.array(dtype=float),
-    contact_pos_out: wp.array(dtype=wp.vec3),
-    contact_frame_out: wp.array(dtype=wp.mat33),
-    contact_includemargin_out: wp.array(dtype=float),
-    contact_friction_out: wp.array(dtype=vec5),
-    contact_solref_out: wp.array(dtype=wp.vec2),
-    contact_solreffriction_out: wp.array(dtype=wp.vec2),
-    contact_solimp_out: wp.array(dtype=vec5),
-    contact_dim_out: wp.array(dtype=int),
-    contact_geom_out: wp.array(dtype=wp.vec2i),
-    contact_efc_address_out: wp.array2d(dtype=int),
-    contact_worldid_out: wp.array(dtype=int),
-    contact_type_out: wp.array(dtype=int),
-    contact_geomcollisionid_out: wp.array(dtype=int),
-    nacon_out: wp.array(dtype=int),
+    contact_dist_out: wp.array[float],
+    contact_pos_out: wp.array[wp.vec3],
+    contact_frame_out: wp.array[wp.mat33],
+    contact_includemargin_out: wp.array[float],
+    contact_friction_out: wp.array[vec5],
+    contact_solref_out: wp.array[wp.vec2],
+    contact_solreffriction_out: wp.array[wp.vec2],
+    contact_solimp_out: wp.array[vec5],
+    contact_dim_out: wp.array[int],
+    contact_geom_out: wp.array[wp.vec2i],
+    contact_efc_address_out: wp.array2d[int],
+    contact_worldid_out: wp.array[int],
+    contact_type_out: wp.array[int],
+    contact_geomcollisionid_out: wp.array[int],
+    nacon_out: wp.array[int],
   ) -> int:
     points = mat43()
     witness1 = mat43()
@@ -892,83 +893,83 @@ def ccd_kernel_builder(
   @wp.kernel(module="unique", enable_backward=False)
   def ccd_kernel(
     # Model:
-    opt_ccd_tolerance: wp.array(dtype=float),
-    geom_type: wp.array(dtype=int),
-    geom_condim: wp.array(dtype=int),
-    geom_dataid: wp.array(dtype=int),
-    geom_priority: wp.array(dtype=int),
-    geom_solmix: wp.array2d(dtype=float),
-    geom_solref: wp.array2d(dtype=wp.vec2),
-    geom_solimp: wp.array2d(dtype=vec5),
-    geom_size: wp.array2d(dtype=wp.vec3),
-    geom_friction: wp.array2d(dtype=wp.vec3),
-    geom_margin: wp.array2d(dtype=float),
-    geom_gap: wp.array2d(dtype=float),
-    mesh_vertadr: wp.array(dtype=int),
-    mesh_vertnum: wp.array(dtype=int),
-    mesh_graphadr: wp.array(dtype=int),
-    mesh_vert: wp.array(dtype=wp.vec3),
-    mesh_graph: wp.array(dtype=int),
-    mesh_polynum: wp.array(dtype=int),
-    mesh_polyadr: wp.array(dtype=int),
-    mesh_polynormal: wp.array(dtype=wp.vec3),
-    mesh_polyvertadr: wp.array(dtype=int),
-    mesh_polyvertnum: wp.array(dtype=int),
-    mesh_polyvert: wp.array(dtype=int),
-    mesh_polymapadr: wp.array(dtype=int),
-    mesh_polymapnum: wp.array(dtype=int),
-    mesh_polymap: wp.array(dtype=int),
-    pair_dim: wp.array(dtype=int),
-    pair_solref: wp.array2d(dtype=wp.vec2),
-    pair_solreffriction: wp.array2d(dtype=wp.vec2),
-    pair_solimp: wp.array2d(dtype=vec5),
-    pair_margin: wp.array2d(dtype=float),
-    pair_gap: wp.array2d(dtype=float),
-    pair_friction: wp.array2d(dtype=vec5),
+    opt_ccd_tolerance: wp.array[float],
+    geom_type: wp.array[int],
+    geom_condim: wp.array[int],
+    geom_dataid: wp.array2d[int],
+    geom_priority: wp.array[int],
+    geom_solmix: wp.array2d[float],
+    geom_solref: wp.array2d[wp.vec2],
+    geom_solimp: wp.array2d[vec5],
+    geom_size: wp.array2d[wp.vec3],
+    geom_friction: wp.array2d[wp.vec3],
+    geom_margin: wp.array2d[float],
+    geom_gap: wp.array2d[float],
+    mesh_vertadr: wp.array[int],
+    mesh_vertnum: wp.array[int],
+    mesh_graphadr: wp.array[int],
+    mesh_vert: wp.array[wp.vec3],
+    mesh_graph: wp.array[int],
+    mesh_polynum: wp.array[int],
+    mesh_polyadr: wp.array[int],
+    mesh_polynormal: wp.array[wp.vec3],
+    mesh_polyvertadr: wp.array[int],
+    mesh_polyvertnum: wp.array[int],
+    mesh_polyvert: wp.array[int],
+    mesh_polymapadr: wp.array[int],
+    mesh_polymapnum: wp.array[int],
+    mesh_polymap: wp.array[int],
+    pair_dim: wp.array[int],
+    pair_solref: wp.array2d[wp.vec2],
+    pair_solreffriction: wp.array2d[wp.vec2],
+    pair_solimp: wp.array2d[vec5],
+    pair_margin: wp.array2d[float],
+    pair_gap: wp.array2d[float],
+    pair_friction: wp.array2d[vec5],
     # Data in:
-    geom_xpos_in: wp.array2d(dtype=wp.vec3),
-    geom_xmat_in: wp.array2d(dtype=wp.mat33),
+    geom_xpos_in: wp.array2d[wp.vec3],
+    geom_xmat_in: wp.array2d[wp.mat33],
     naconmax_in: int,
     naccdmax_in: int,
-    ncollision_in: wp.array(dtype=int),
+    ncollision_in: wp.array[int],
     # In:
-    collision_pair_in: wp.array(dtype=wp.vec2i),
-    collision_pairid_in: wp.array(dtype=wp.vec2i),
-    collision_worldid_in: wp.array(dtype=int),
-    epa_vert_in: wp.array2d(dtype=wp.vec3),
-    epa_vert_index_in: wp.array2d(dtype=int),
-    epa_face_in: wp.array2d(dtype=int),
-    epa_pr_in: wp.array2d(dtype=wp.vec3),
-    epa_norm2_in: wp.array2d(dtype=float),
-    epa_horizon_in: wp.array2d(dtype=int),
-    multiccd_polygon_in: wp.array2d(dtype=wp.vec3),
-    multiccd_clipped_in: wp.array2d(dtype=wp.vec3),
-    multiccd_pnormal_in: wp.array2d(dtype=wp.vec3),
-    multiccd_pdist_in: wp.array2d(dtype=float),
-    multiccd_idx1_in: wp.array2d(dtype=int),
-    multiccd_idx2_in: wp.array2d(dtype=int),
-    multiccd_n1_in: wp.array2d(dtype=wp.vec3),
-    multiccd_n2_in: wp.array2d(dtype=wp.vec3),
-    multiccd_endvert_in: wp.array2d(dtype=wp.vec3),
-    multiccd_face1_in: wp.array2d(dtype=wp.vec3),
-    multiccd_face2_in: wp.array2d(dtype=wp.vec3),
-    nccd_in: wp.array(dtype=int),
+    collision_pair_in: wp.array[wp.vec2i],
+    collision_pairid_in: wp.array[wp.vec2i],
+    collision_worldid_in: wp.array[int],
+    epa_vert_in: wp.array2d[wp.vec3],
+    epa_vert_index_in: wp.array2d[int],
+    epa_face_in: wp.array2d[int],
+    epa_pr_in: wp.array2d[wp.vec3],
+    epa_norm2_in: wp.array2d[float],
+    epa_horizon_in: wp.array2d[int],
+    multiccd_polygon_in: wp.array2d[wp.vec3],
+    multiccd_clipped_in: wp.array2d[wp.vec3],
+    multiccd_pnormal_in: wp.array2d[wp.vec3],
+    multiccd_pdist_in: wp.array2d[float],
+    multiccd_idx1_in: wp.array2d[int],
+    multiccd_idx2_in: wp.array2d[int],
+    multiccd_n1_in: wp.array2d[wp.vec3],
+    multiccd_n2_in: wp.array2d[wp.vec3],
+    multiccd_endvert_in: wp.array2d[wp.vec3],
+    multiccd_face1_in: wp.array2d[wp.vec3],
+    multiccd_face2_in: wp.array2d[wp.vec3],
+    nccd_in: wp.array[int],
     # Data out:
-    contact_dist_out: wp.array(dtype=float),
-    contact_pos_out: wp.array(dtype=wp.vec3),
-    contact_frame_out: wp.array(dtype=wp.mat33),
-    contact_includemargin_out: wp.array(dtype=float),
-    contact_friction_out: wp.array(dtype=vec5),
-    contact_solref_out: wp.array(dtype=wp.vec2),
-    contact_solreffriction_out: wp.array(dtype=wp.vec2),
-    contact_solimp_out: wp.array(dtype=vec5),
-    contact_dim_out: wp.array(dtype=int),
-    contact_geom_out: wp.array(dtype=wp.vec2i),
-    contact_efc_address_out: wp.array2d(dtype=int),
-    contact_worldid_out: wp.array(dtype=int),
-    contact_type_out: wp.array(dtype=int),
-    contact_geomcollisionid_out: wp.array(dtype=int),
-    nacon_out: wp.array(dtype=int),
+    contact_dist_out: wp.array[float],
+    contact_pos_out: wp.array[wp.vec3],
+    contact_frame_out: wp.array[wp.mat33],
+    contact_includemargin_out: wp.array[float],
+    contact_friction_out: wp.array[vec5],
+    contact_solref_out: wp.array[wp.vec2],
+    contact_solreffriction_out: wp.array[wp.vec2],
+    contact_solimp_out: wp.array[vec5],
+    contact_dim_out: wp.array[int],
+    contact_geom_out: wp.array[wp.vec2i],
+    contact_efc_address_out: wp.array2d[int],
+    contact_worldid_out: wp.array[int],
+    contact_type_out: wp.array[int],
+    contact_geomcollisionid_out: wp.array[int],
+    nacon_out: wp.array[int],
   ):
     collisionid = wp.tid()
     if collisionid >= ncollision_in[0]:

mjx/mujoco/mjx/third_party/mujoco_warp/_src/collision_core.py

@@ -47,44 +47,44 @@ class Geom:
   hfprism: mat63
   vertadr: int
   vertnum: int
-  vert: wp.array(dtype=wp.vec3)
+  vert: wp.array[wp.vec3]
   graphadr: int
-  graph: wp.array(dtype=int)
+  graph: wp.array[int]
   mesh_polynum: int
   mesh_polyadr: int
-  mesh_polynormal: wp.array(dtype=wp.vec3)
-  mesh_polyvertadr: wp.array(dtype=int)
-  mesh_polyvertnum: wp.array(dtype=int)
-  mesh_polyvert: wp.array(dtype=int)
-  mesh_polymapadr: wp.array(dtype=int)
-  mesh_polymapnum: wp.array(dtype=int)
-  mesh_polymap: wp.array(dtype=int)
+  mesh_polynormal: wp.array[wp.vec3]
+  mesh_polyvertadr: wp.array[int]
+  mesh_polyvertnum: wp.array[int]
+  mesh_polyvert: wp.array[int]
+  mesh_polymapadr: wp.array[int]
+  mesh_polymapnum: wp.array[int]
+  mesh_polymap: wp.array[int]
   index: int
 
 
 @wp.func
 def geom_collision_pair(
   # Model:
-  geom_type: wp.array(dtype=int),
-  geom_dataid: wp.array(dtype=int),
-  geom_size: wp.array2d(dtype=wp.vec3),
-  mesh_vertadr: wp.array(dtype=int),
-  mesh_vertnum: wp.array(dtype=int),
-  mesh_graphadr: wp.array(dtype=int),
-  mesh_vert: wp.array(dtype=wp.vec3),
-  mesh_graph: wp.array(dtype=int),
-  mesh_polynum: wp.array(dtype=int),
-  mesh_polyadr: wp.array(dtype=int),
-  mesh_polynormal: wp.array(dtype=wp.vec3),
-  mesh_polyvertadr: wp.array(dtype=int),
-  mesh_polyvertnum: wp.array(dtype=int),
-  mesh_polyvert: wp.array(dtype=int),
-  mesh_polymapadr: wp.array(dtype=int),
-  mesh_polymapnum: wp.array(dtype=int),
-  mesh_polymap: wp.array(dtype=int),
+  geom_type: wp.array[int],
+  geom_dataid: wp.array2d[int],
+  geom_size: wp.array2d[wp.vec3],
+  mesh_vertadr: wp.array[int],
+  mesh_vertnum: wp.array[int],
+  mesh_graphadr: wp.array[int],
+  mesh_vert: wp.array[wp.vec3],
+  mesh_graph: wp.array[int],
+  mesh_polynum: wp.array[int],
+  mesh_polyadr: wp.array[int],
+  mesh_polynormal: wp.array[wp.vec3],
+  mesh_polyvertadr: wp.array[int],
+  mesh_polyvertnum: wp.array[int],
+  mesh_polyvert: wp.array[int],
+  mesh_polymapadr: wp.array[int],
+  mesh_polymapnum: wp.array[int],
+  mesh_polymap: wp.array[int],
   # Data in:
-  geom_xpos_in: wp.array2d(dtype=wp.vec3),
-  geom_xmat_in: wp.array2d(dtype=wp.mat33),
+  geom_xpos_in: wp.array2d[wp.vec3],
+  geom_xmat_in: wp.array2d[wp.mat33],
   # In:
   geoms: wp.vec2i,
   worldid: int,
@@ -109,8 +109,10 @@ def geom_collision_pair(
   # z-axis of the rotation matrix, used as the surface normal for plane collisions
   geom2.normal = wp.vec3(geom2.rot[0, 2], geom2.rot[1, 2], geom2.rot[2, 2])
 
+  dataid_setid = worldid % geom_dataid.shape[0]
+
   if geom_type1 == GeomType.MESH:
-    dataid = geom_dataid[g1]
+    dataid = geom_dataid[dataid_setid, g1]
     geom1.vertadr = wp.where(dataid >= 0, mesh_vertadr[dataid], -1)
     geom1.vertnum = wp.where(dataid >= 0, mesh_vertnum[dataid], -1)
     geom1.graphadr = wp.where(dataid >= 0, mesh_graphadr[dataid], -1)
@@ -128,7 +130,7 @@ def geom_collision_pair(
     geom1.mesh_polymap = mesh_polymap
 
   if geom_type2 == GeomType.MESH:
-    dataid = geom_dataid[g2]
+    dataid = geom_dataid[dataid_setid, g2]
     geom2.vertadr = wp.where(dataid >= 0, mesh_vertadr[dataid], -1)
     geom2.vertnum = wp.where(dataid >= 0, mesh_vertnum[dataid], -1)
     geom2.graphadr = wp.where(dataid >= 0, mesh_graphadr[dataid], -1)
@@ -174,21 +176,21 @@ def write_contact(
   pairid_in: wp.vec2i,
   worldid_in: int,
   # Data out:
-  contact_dist_out: wp.array(dtype=float),
-  contact_pos_out: wp.array(dtype=wp.vec3),
-  contact_frame_out: wp.array(dtype=wp.mat33),
-  contact_includemargin_out: wp.array(dtype=float),
-  contact_friction_out: wp.array(dtype=vec5),
-  contact_solref_out: wp.array(dtype=wp.vec2),
-  contact_solreffriction_out: wp.array(dtype=wp.vec2),
-  contact_solimp_out: wp.array(dtype=vec5),
-  contact_dim_out: wp.array(dtype=int),
-  contact_geom_out: wp.array(dtype=wp.vec2i),
-  contact_efc_address_out: wp.array2d(dtype=int),
-  contact_worldid_out: wp.array(dtype=int),
-  contact_type_out: wp.array(dtype=int),
-  contact_geomcollisionid_out: wp.array(dtype=int),
-  nacon_out: wp.array(dtype=int),
+  contact_dist_out: wp.array[float],
+  contact_pos_out: wp.array[wp.vec3],
+  contact_frame_out: wp.array[wp.mat33],
+  contact_includemargin_out: wp.array[float],
+  contact_friction_out: wp.array[vec5],
+  contact_solref_out: wp.array[wp.vec2],
+  contact_solreffriction_out: wp.array[wp.vec2],
+  contact_solimp_out: wp.array[vec5],
+  contact_dim_out: wp.array[int],
+  contact_geom_out: wp.array[wp.vec2i],
+  contact_efc_address_out: wp.array2d[int],
+  contact_worldid_out: wp.array[int],
+  contact_type_out: wp.array[int],
+  contact_geomcollisionid_out: wp.array[int],
+  nacon_out: wp.array[int],
 ) -> int:
   """Atomically write a detected contact into the contact output arrays.
 
@@ -233,24 +235,24 @@ def write_contact(
 @wp.func
 def contact_params(
   # Model:
-  geom_condim: wp.array(dtype=int),
-  geom_priority: wp.array(dtype=int),
-  geom_solmix: wp.array2d(dtype=float),
-  geom_solref: wp.array2d(dtype=wp.vec2),
-  geom_solimp: wp.array2d(dtype=vec5),
-  geom_friction: wp.array2d(dtype=wp.vec3),
-  geom_margin: wp.array2d(dtype=float),
-  geom_gap: wp.array2d(dtype=float),
-  pair_dim: wp.array(dtype=int),
-  pair_solref: wp.array2d(dtype=wp.vec2),
-  pair_solreffriction: wp.array2d(dtype=wp.vec2),
-  pair_solimp: wp.array2d(dtype=vec5),
-  pair_margin: wp.array2d(dtype=float),
-  pair_gap: wp.array2d(dtype=float),
-  pair_friction: wp.array2d(dtype=vec5),
+  geom_condim: wp.array[int],
+  geom_priority: wp.array[int],
+  geom_solmix: wp.array2d[float],
+  geom_solref: wp.array2d[wp.vec2],
+  geom_solimp: wp.array2d[vec5],
+  geom_friction: wp.array2d[wp.vec3],
+  geom_margin: wp.array2d[float],
+  geom_gap: wp.array2d[float],
+  pair_dim: wp.array[int],
+  pair_solref: wp.array2d[wp.vec2],
+  pair_solreffriction: wp.array2d[wp.vec2],
+  pair_solimp: wp.array2d[vec5],
+  pair_margin: wp.array2d[float],
+  pair_gap: wp.array2d[float],
+  pair_friction: wp.array2d[vec5],
   # In:
-  collision_pair_in: wp.array(dtype=wp.vec2i),
-  collision_pairid_in: wp.array(dtype=wp.vec2i),
+  collision_pair_in: wp.array[wp.vec2i],
+  collision_pairid_in: wp.array[wp.vec2i],
   cid: int,
   worldid: int,
 ):

mjx/mujoco/mjx/third_party/mujoco_warp/_src/collision_driver.py

@@ -80,8 +80,8 @@ MJ_COLLISION_TABLE = {
 @wp.kernel
 def _zero_nacon_ncollision(
   # Data out:
-  nacon_out: wp.array(dtype=int),
-  ncollision_out: wp.array(dtype=int),
+  nacon_out: wp.array[int],
+  ncollision_out: wp.array[int],
 ):
   ncollision_out[0] = 0
   nacon_out[0] = 0
@@ -275,12 +275,12 @@ def _broadphase_filter(opt_broadphase_filter: int, ngeom_aabb: int, ngeom_rbound
   @wp.func
   def func(
     # Model:
-    geom_aabb: wp.array3d(dtype=wp.vec3),
-    geom_rbound: wp.array2d(dtype=float),
-    geom_margin: wp.array2d(dtype=float),
+    geom_aabb: wp.array3d[wp.vec3],
+    geom_rbound: wp.array2d[float],
+    geom_margin: wp.array2d[float],
     # Data in:
-    geom_xpos_in: wp.array2d(dtype=wp.vec3),
-    geom_xmat_in: wp.array2d(dtype=wp.mat33),
+    geom_xpos_in: wp.array2d[wp.vec3],
+    geom_xmat_in: wp.array2d[wp.mat33],
     # In:
     geom1: int,
     geom2: int,
@@ -324,8 +324,8 @@ def _broadphase_filter(opt_broadphase_filter: int, ngeom_aabb: int, ngeom_rbound
 @wp.func
 def _add_geom_pair(
   # Model:
-  geom_type: wp.array(dtype=int),
-  nxn_pairid: wp.array(dtype=wp.vec2i),
+  geom_type: wp.array[int],
+  nxn_pairid: wp.array[wp.vec2i],
   # Data in:
   naconmax_in: int,
   # In:
@@ -334,11 +334,11 @@ def _add_geom_pair(
   worldid: int,
   nxnid: int,
   # Data out:
-  ncollision_out: wp.array(dtype=int),
+  ncollision_out: wp.array[int],
   # Out:
-  collision_pair_out: wp.array(dtype=wp.vec2i),
-  collision_pairid_out: wp.array(dtype=wp.vec2i),
-  collision_worldid_out: wp.array(dtype=int),
+  collision_pair_out: wp.array[wp.vec2i],
+  collision_pairid_out: wp.array[wp.vec2i],
+  collision_worldid_out: wp.array[int],
 ):
   pairid = wp.atomic_add(ncollision_out, 0, 1)
 
@@ -359,7 +359,7 @@ def _add_geom_pair(
 
 
 @wp.func
-def _binary_search(values: wp.array(dtype=Any), value: Any, lower: int, upper: int) -> int:
+def _binary_search(values: wp.array[Any], value: Any, lower: int, upper: int) -> int:
   while lower < upper:
     mid = (lower + upper) >> 1
     if values[mid] > value:
@@ -375,18 +375,18 @@ def _sap_project(opt_broadphase: int):
   def sap_project(
     # Model:
     ngeom: int,
-    geom_rbound: wp.array2d(dtype=float),
-    geom_margin: wp.array2d(dtype=float),
+    geom_rbound: wp.array2d[float],
+    geom_margin: wp.array2d[float],
     # Data in:
-    geom_xpos_in: wp.array2d(dtype=wp.vec3),
+    geom_xpos_in: wp.array2d[wp.vec3],
     nworld_in: int,
     # In:
     direction_in: wp.vec3,
     # Out:
-    projection_lower_out: wp.array2d(dtype=float),
-    projection_upper_out: wp.array2d(dtype=float),
-    sort_index_out: wp.array2d(dtype=int),
-    segmented_index_out: wp.array(dtype=int),
+    projection_lower_out: wp.array2d[float],
+    projection_upper_out: wp.array2d[float],
+    sort_index_out: wp.array2d[int],
+    segmented_index_out: wp.array[int],
   ):
     worldid, geomid = wp.tid()
 
@@ -422,11 +422,11 @@ def _sap_range(
   # Model:
   ngeom: int,
   # In:
-  projection_lower_in: wp.array2d(dtype=float),
-  projection_upper_in: wp.array2d(dtype=float),
-  sort_index_in: wp.array2d(dtype=int),
+  projection_lower_in: wp.array2d[float],
+  projection_upper_in: wp.array2d[float],
+  sort_index_in: wp.array2d[int],
   # Out:
-  range_out: wp.array2d(dtype=int),
+  range_out: wp.array2d[int],
 ):
   worldid, geomid = wp.tid()
 
@@ -448,26 +448,26 @@ def _sap_broadphase(opt_broadphase_filter: int, ngeom_aabb: int, ngeom_rbound: i
   def kernel(
     # Model:
     ngeom: int,
-    geom_type: wp.array(dtype=int),
-    geom_aabb: wp.array3d(dtype=wp.vec3),
-    geom_rbound: wp.array2d(dtype=float),
-    geom_margin: wp.array2d(dtype=float),
-    nxn_pairid: wp.array(dtype=wp.vec2i),
+    geom_type: wp.array[int],
+    geom_aabb: wp.array3d[wp.vec3],
+    geom_rbound: wp.array2d[float],
+    geom_margin: wp.array2d[float],
+    nxn_pairid: wp.array[wp.vec2i],
     # Data in:
-    geom_xpos_in: wp.array2d(dtype=wp.vec3),
-    geom_xmat_in: wp.array2d(dtype=wp.mat33),
+    geom_xpos_in: wp.array2d[wp.vec3],
+    geom_xmat_in: wp.array2d[wp.mat33],
     nworld_in: int,
     naconmax_in: int,
     # In:
-    sort_index_in: wp.array2d(dtype=int),
-    cumulative_sum_in: wp.array(dtype=int),
+    sort_index_in: wp.array2d[int],
+    cumulative_sum_in: wp.array[int],
     nsweep_in: int,
     # Data out:
-    ncollision_out: wp.array(dtype=int),
+    ncollision_out: wp.array[int],
     # Out:
-    collision_pair_out: wp.array(dtype=wp.vec2i),
-    collision_pairid_out: wp.array(dtype=wp.vec2i),
-    collision_worldid_out: wp.array(dtype=int),
+    collision_pair_out: wp.array[wp.vec2i],
+    collision_pairid_out: wp.array[wp.vec2i],
+    collision_worldid_out: wp.array[int],
   ):
     worldgeomid = wp.tid()
 
@@ -528,11 +528,11 @@ def _segmented_sort(tile_size: int):
   @wp.kernel(module="unique")
   def segmented_sort(
     # In:
-    projection_lower_in: wp.array2d(dtype=float),
-    sort_index_in: wp.array2d(dtype=int),
+    projection_lower_in: wp.array2d[float],
+    sort_index_in: wp.array2d[int],
     # Out:
-    projection_lower_out: wp.array2d(dtype=float),
-    sort_index_out: wp.array2d(dtype=int),
+    projection_lower_out: wp.array2d[float],
+    sort_index_out: wp.array2d[int],
   ):
     worldid = wp.tid()
 
@@ -648,22 +648,22 @@ def _nxn_broadphase(opt_broadphase_filter: int, ngeom_aabb: int, ngeom_rbound: i
   @wp.kernel(module="unique", enable_backward=False)
   def kernel(
     # Model:
-    geom_type: wp.array(dtype=int),
-    geom_aabb: wp.array3d(dtype=wp.vec3),
-    geom_rbound: wp.array2d(dtype=float),
-    geom_margin: wp.array2d(dtype=float),
-    nxn_geom_pair: wp.array(dtype=wp.vec2i),
-    nxn_pairid: wp.array(dtype=wp.vec2i),
+    geom_type: wp.array[int],
+    geom_aabb: wp.array3d[wp.vec3],
+    geom_rbound: wp.array2d[float],
+    geom_margin: wp.array2d[float],
+    nxn_geom_pair: wp.array[wp.vec2i],
+    nxn_pairid: wp.array[wp.vec2i],
     # Data in:
-    geom_xpos_in: wp.array2d(dtype=wp.vec3),
-    geom_xmat_in: wp.array2d(dtype=wp.mat33),
+    geom_xpos_in: wp.array2d[wp.vec3],
+    geom_xmat_in: wp.array2d[wp.mat33],
     naconmax_in: int,
     # Data out:
-    ncollision_out: wp.array(dtype=int),
+    ncollision_out: wp.array[int],
     # Out:
-    collision_pair_out: wp.array(dtype=wp.vec2i),
-    collision_pairid_out: wp.array(dtype=wp.vec2i),
-    collision_worldid_out: wp.array(dtype=int),
+    collision_pair_out: wp.array[wp.vec2i],
+    collision_pairid_out: wp.array[wp.vec2i],
+    collision_worldid_out: wp.array[int],
   ):
     worldid, elementid = wp.tid()
 

mjx/mujoco/mjx/third_party/mujoco_warp/_src/collision_flex.py

@@ -46,22 +46,22 @@ def _write_flex_contact(
   vertid: int,
   worldid: int,
   # Data out:
-  contact_dist_out: wp.array(dtype=float),
-  contact_pos_out: wp.array(dtype=wp.vec3),
-  contact_frame_out: wp.array(dtype=wp.mat33),
-  contact_includemargin_out: wp.array(dtype=float),
-  contact_friction_out: wp.array(dtype=vec5),
-  contact_solref_out: wp.array(dtype=wp.vec2),
-  contact_solreffriction_out: wp.array(dtype=wp.vec2),
-  contact_solimp_out: wp.array(dtype=vec5),
-  contact_dim_out: wp.array(dtype=int),
-  contact_geom_out: wp.array(dtype=wp.vec2i),
-  contact_flex_out: wp.array(dtype=wp.vec2i),
-  contact_vert_out: wp.array(dtype=wp.vec2i),
-  contact_worldid_out: wp.array(dtype=int),
-  contact_type_out: wp.array(dtype=int),
-  contact_geomcollisionid_out: wp.array(dtype=int),
-  nacon_out: wp.array(dtype=int),
+  contact_dist_out: wp.array[float],
+  contact_pos_out: wp.array[wp.vec3],
+  contact_frame_out: wp.array[wp.mat33],
+  contact_includemargin_out: wp.array[float],
+  contact_friction_out: wp.array[vec5],
+  contact_solref_out: wp.array[wp.vec2],
+  contact_solreffriction_out: wp.array[wp.vec2],
+  contact_solimp_out: wp.array[vec5],
+  contact_dim_out: wp.array[int],
+  contact_geom_out: wp.array[wp.vec2i],
+  contact_flex_out: wp.array[wp.vec2i],
+  contact_vert_out: wp.array[wp.vec2i],
+  contact_worldid_out: wp.array[int],
+  contact_type_out: wp.array[int],
+  contact_geomcollisionid_out: wp.array[int],
+  nacon_out: wp.array[int],
 ):
   if dist >= margin or dist >= MJ_MAXVAL:
     return
@@ -110,22 +110,22 @@ def _collide_geom_triangle(
   vertex_id: int,
   worldid: int,
   # Data out:
-  contact_dist_out: wp.array(dtype=float),
-  contact_pos_out: wp.array(dtype=wp.vec3),
-  contact_frame_out: wp.array(dtype=wp.mat33),
-  contact_includemargin_out: wp.array(dtype=float),
-  contact_friction_out: wp.array(dtype=vec5),
-  contact_solref_out: wp.array(dtype=wp.vec2),
-  contact_solreffriction_out: wp.array(dtype=wp.vec2),
-  contact_solimp_out: wp.array(dtype=vec5),
-  contact_dim_out: wp.array(dtype=int),
-  contact_geom_out: wp.array(dtype=wp.vec2i),
-  contact_flex_out: wp.array(dtype=wp.vec2i),
-  contact_vert_out: wp.array(dtype=wp.vec2i),
-  contact_worldid_out: wp.array(dtype=int),
-  contact_type_out: wp.array(dtype=int),
-  contact_geomcollisionid_out: wp.array(dtype=int),
-  nacon_out: wp.array(dtype=int),
+  contact_dist_out: wp.array[float],
+  contact_pos_out: wp.array[wp.vec3],
+  contact_frame_out: wp.array[wp.mat33],
+  contact_includemargin_out: wp.array[float],
+  contact_friction_out: wp.array[vec5],
+  contact_solref_out: wp.array[wp.vec2],
+  contact_solreffriction_out: wp.array[wp.vec2],
+  contact_solimp_out: wp.array[vec5],
+  contact_dim_out: wp.array[int],
+  contact_geom_out: wp.array[wp.vec2i],
+  contact_flex_out: wp.array[wp.vec2i],
+  contact_vert_out: wp.array[wp.vec2i],
+  contact_worldid_out: wp.array[int],
+  contact_type_out: wp.array[int],
+  contact_geomcollisionid_out: wp.array[int],
+  nacon_out: wp.array[int],
 ):
   if gtype == int(GeomType.SPHERE):
     sphere_radius = size_val[0]
@@ -262,41 +262,41 @@ def _flex_plane_narrowphase(
   # Model:
   ngeom: int,
   nflexvert: int,
-  geom_type: wp.array(dtype=int),
-  geom_condim: wp.array(dtype=int),
-  geom_solref: wp.array2d(dtype=wp.vec2),
-  geom_solimp: wp.array2d(dtype=vec5),
-  geom_friction: wp.array2d(dtype=wp.vec3),
-  geom_margin: wp.array2d(dtype=float),
-  flex_condim: wp.array(dtype=int),
-  flex_friction: wp.array(dtype=wp.vec3),
-  flex_margin: wp.array(dtype=float),
-  flex_vertadr: wp.array(dtype=int),
-  flex_radius: wp.array(dtype=float),
-  flex_vertflexid: wp.array(dtype=int),
+  geom_type: wp.array[int],
+  geom_condim: wp.array[int],
+  geom_solref: wp.array2d[wp.vec2],
+  geom_solimp: wp.array2d[vec5],
+  geom_friction: wp.array2d[wp.vec3],
+  geom_margin: wp.array2d[float],
+  flex_condim: wp.array[int],
+  flex_friction: wp.array[wp.vec3],
+  flex_margin: wp.array[float],
+  flex_vertadr: wp.array[int],
+  flex_radius: wp.array[float],
+  flex_vertflexid: wp.array[int],
   # Data in:
-  geom_xpos_in: wp.array2d(dtype=wp.vec3),
-  geom_xmat_in: wp.array2d(dtype=wp.mat33),
-  flexvert_xpos_in: wp.array2d(dtype=wp.vec3),
+  geom_xpos_in: wp.array2d[wp.vec3],
+  geom_xmat_in: wp.array2d[wp.mat33],
+  flexvert_xpos_in: wp.array2d[wp.vec3],
   nworld_in: int,
   naconmax_in: int,
   # Data out:
-  contact_dist_out: wp.array(dtype=float),
-  contact_pos_out: wp.array(dtype=wp.vec3),
-  contact_frame_out: wp.array(dtype=wp.mat33),
-  contact_includemargin_out: wp.array(dtype=float),
-  contact_friction_out: wp.array(dtype=vec5),
-  contact_solref_out: wp.array(dtype=wp.vec2),
-  contact_solreffriction_out: wp.array(dtype=wp.vec2),
-  contact_solimp_out: wp.array(dtype=vec5),
-  contact_dim_out: wp.array(dtype=int),
-  contact_geom_out: wp.array(dtype=wp.vec2i),
-  contact_flex_out: wp.array(dtype=wp.vec2i),
-  contact_vert_out: wp.array(dtype=wp.vec2i),
-  contact_worldid_out: wp.array(dtype=int),
-  contact_type_out: wp.array(dtype=int),
-  contact_geomcollisionid_out: wp.array(dtype=int),
-  nacon_out: wp.array(dtype=int),
+  contact_dist_out: wp.array[float],
+  contact_pos_out: wp.array[wp.vec3],
+  contact_frame_out: wp.array[wp.mat33],
+  contact_includemargin_out: wp.array[float],
+  contact_friction_out: wp.array[vec5],
+  contact_solref_out: wp.array[wp.vec2],
+  contact_solreffriction_out: wp.array[wp.vec2],
+  contact_solimp_out: wp.array[vec5],
+  contact_dim_out: wp.array[int],
+  contact_geom_out: wp.array[wp.vec2i],
+  contact_flex_out: wp.array[wp.vec2i],
+  contact_vert_out: wp.array[wp.vec2i],
+  contact_worldid_out: wp.array[int],
+  contact_type_out: wp.array[int],
+  contact_geomcollisionid_out: wp.array[int],
+  nacon_out: wp.array[int],
 ):
   worldid, vertid = wp.tid()
 
@@ -382,48 +382,48 @@ def _flex_narrowphase_dim2(
   # Model:
   ngeom: int,
   nflex: int,
-  geom_type: wp.array(dtype=int),
-  geom_contype: wp.array(dtype=int),
-  geom_conaffinity: wp.array(dtype=int),
-  geom_condim: wp.array(dtype=int),
-  geom_solref: wp.array2d(dtype=wp.vec2),
-  geom_solimp: wp.array2d(dtype=vec5),
-  geom_size: wp.array2d(dtype=wp.vec3),
-  geom_friction: wp.array2d(dtype=wp.vec3),
-  geom_margin: wp.array2d(dtype=float),
-  flex_contype: wp.array(dtype=int),
-  flex_conaffinity: wp.array(dtype=int),
-  flex_margin: wp.array(dtype=float),
-  flex_dim: wp.array(dtype=int),
-  flex_vertadr: wp.array(dtype=int),
-  flex_elemadr: wp.array(dtype=int),
-  flex_elemnum: wp.array(dtype=int),
-  flex_elemdataadr: wp.array(dtype=int),
-  flex_elem: wp.array(dtype=int),
-  flex_radius: wp.array(dtype=float),
+  geom_type: wp.array[int],
+  geom_contype: wp.array[int],
+  geom_conaffinity: wp.array[int],
+  geom_condim: wp.array[int],
+  geom_solref: wp.array2d[wp.vec2],
+  geom_solimp: wp.array2d[vec5],
+  geom_size: wp.array2d[wp.vec3],
+  geom_friction: wp.array2d[wp.vec3],
+  geom_margin: wp.array2d[float],
+  flex_contype: wp.array[int],
+  flex_conaffinity: wp.array[int],
+  flex_margin: wp.array[float],
+  flex_dim: wp.array[int],
+  flex_vertadr: wp.array[int],
+  flex_elemadr: wp.array[int],
+  flex_elemnum: wp.array[int],
+  flex_elemdataadr: wp.array[int],
+  flex_elem: wp.array[int],
+  flex_radius: wp.array[float],
   # Data in:
-  geom_xpos_in: wp.array2d(dtype=wp.vec3),
-  geom_xmat_in: wp.array2d(dtype=wp.mat33),
-  flexvert_xpos_in: wp.array2d(dtype=wp.vec3),
+  geom_xpos_in: wp.array2d[wp.vec3],
+  geom_xmat_in: wp.array2d[wp.mat33],
+  flexvert_xpos_in: wp.array2d[wp.vec3],
   nworld_in: int,
   naconmax_in: int,
   # Data out:
-  contact_dist_out: wp.array(dtype=float),
-  contact_pos_out: wp.array(dtype=wp.vec3),
-  contact_frame_out: wp.array(dtype=wp.mat33),
-  contact_includemargin_out: wp.array(dtype=float),
-  contact_friction_out: wp.array(dtype=vec5),
-  contact_solref_out: wp.array(dtype=wp.vec2),
-  contact_solreffriction_out: wp.array(dtype=wp.vec2),
-  contact_solimp_out: wp.array(dtype=vec5),
-  contact_dim_out: wp.array(dtype=int),
-  contact_geom_out: wp.array(dtype=wp.vec2i),
-  contact_flex_out: wp.array(dtype=wp.vec2i),
-  contact_vert_out: wp.array(dtype=wp.vec2i),
-  contact_worldid_out: wp.array(dtype=int),
-  contact_type_out: wp.array(dtype=int),
-  contact_geomcollisionid_out: wp.array(dtype=int),
-  nacon_out: wp.array(dtype=int),
+  contact_dist_out: wp.array[float],
+  contact_pos_out: wp.array[wp.vec3],
+  contact_frame_out: wp.array[wp.mat33],
+  contact_includemargin_out: wp.array[float],
+  contact_friction_out: wp.array[vec5],
+  contact_solref_out: wp.array[wp.vec2],
+  contact_solreffriction_out: wp.array[wp.vec2],
+  contact_solimp_out: wp.array[vec5],
+  contact_dim_out: wp.array[int],
+  contact_geom_out: wp.array[wp.vec2i],
+  contact_flex_out: wp.array[wp.vec2i],
+  contact_vert_out: wp.array[wp.vec2i],
+  contact_worldid_out: wp.array[int],
+  contact_type_out: wp.array[int],
+  contact_geomcollisionid_out: wp.array[int],
+  nacon_out: wp.array[int],
 ):
   worldid, elemid = wp.tid()
 
@@ -533,47 +533,47 @@ def _flex_narrowphase_dim3(
   # Model:
   ngeom: int,
   nflex: int,
-  geom_type: wp.array(dtype=int),
-  geom_contype: wp.array(dtype=int),
-  geom_conaffinity: wp.array(dtype=int),
-  geom_condim: wp.array(dtype=int),
-  geom_solref: wp.array2d(dtype=wp.vec2),
-  geom_solimp: wp.array2d(dtype=vec5),
-  geom_size: wp.array2d(dtype=wp.vec3),
-  geom_friction: wp.array2d(dtype=wp.vec3),
-  geom_margin: wp.array2d(dtype=float),
-  flex_contype: wp.array(dtype=int),
-  flex_conaffinity: wp.array(dtype=int),
-  flex_margin: wp.array(dtype=float),
-  flex_dim: wp.array(dtype=int),
-  flex_vertadr: wp.array(dtype=int),
-  flex_shellnum: wp.array(dtype=int),
-  flex_shelldataadr: wp.array(dtype=int),
-  flex_shell: wp.array(dtype=int),
-  flex_radius: wp.array(dtype=float),
+  geom_type: wp.array[int],
+  geom_contype: wp.array[int],
+  geom_conaffinity: wp.array[int],
+  geom_condim: wp.array[int],
+  geom_solref: wp.array2d[wp.vec2],
+  geom_solimp: wp.array2d[vec5],
+  geom_size: wp.array2d[wp.vec3],
+  geom_friction: wp.array2d[wp.vec3],
+  geom_margin: wp.array2d[float],
+  flex_contype: wp.array[int],
+  flex_conaffinity: wp.array[int],
+  flex_margin: wp.array[float],
+  flex_dim: wp.array[int],
+  flex_vertadr: wp.array[int],
+  flex_shellnum: wp.array[int],
+  flex_shelldataadr: wp.array[int],
+  flex_shell: wp.array[int],
+  flex_radius: wp.array[float],
   # Data in:
-  geom_xpos_in: wp.array2d(dtype=wp.vec3),
-  geom_xmat_in: wp.array2d(dtype=wp.mat33),
-  flexvert_xpos_in: wp.array2d(dtype=wp.vec3),
+  geom_xpos_in: wp.array2d[wp.vec3],
+  geom_xmat_in: wp.array2d[wp.mat33],
+  flexvert_xpos_in: wp.array2d[wp.vec3],
   nworld_in: int,
   naconmax_in: int,
   # Data out:
-  contact_dist_out: wp.array(dtype=float),
-  contact_pos_out: wp.array(dtype=wp.vec3),
-  contact_frame_out: wp.array(dtype=wp.mat33),
-  contact_includemargin_out: wp.array(dtype=float),
-  contact_friction_out: wp.array(dtype=vec5),
-  contact_solref_out: wp.array(dtype=wp.vec2),
-  contact_solreffriction_out: wp.array(dtype=wp.vec2),
-  contact_solimp_out: wp.array(dtype=vec5),
-  contact_dim_out: wp.array(dtype=int),
-  contact_geom_out: wp.array(dtype=wp.vec2i),
-  contact_flex_out: wp.array(dtype=wp.vec2i),
-  contact_vert_out: wp.array(dtype=wp.vec2i),
-  contact_worldid_out: wp.array(dtype=int),
-  contact_type_out: wp.array(dtype=int),
-  contact_geomcollisionid_out: wp.array(dtype=int),
-  nacon_out: wp.array(dtype=int),
+  contact_dist_out: wp.array[float],
+  contact_pos_out: wp.array[wp.vec3],
+  contact_frame_out: wp.array[wp.mat33],
+  contact_includemargin_out: wp.array[float],
+  contact_friction_out: wp.array[vec5],
+  contact_solref_out: wp.array[wp.vec2],
+  contact_solreffriction_out: wp.array[wp.vec2],
+  contact_solimp_out: wp.array[vec5],
+  contact_dim_out: wp.array[int],
+  contact_geom_out: wp.array[wp.vec2i],
+  contact_flex_out: wp.array[wp.vec2i],
+  contact_vert_out: wp.array[wp.vec2i],
+  contact_worldid_out: wp.array[int],
+  contact_type_out: wp.array[int],
+  contact_geomcollisionid_out: wp.array[int],
+  nacon_out: wp.array[int],
 ):
   worldid, shellid = wp.tid()
 

mjx/mujoco/mjx/third_party/mujoco_warp/_src/collision_gjk.py

@@ -63,20 +63,20 @@ class Polytope:
   status: int
 
   # vertices in polytope (packed geom1 followed by geom2)
-  vert: wp.array(dtype=wp.vec3)
-  vert_index: wp.array(dtype=int)
+  vert: wp.array[wp.vec3]
+  vert_index: wp.array[int]
   nvert: int
 
   # faces in polytope
   # 10 bits per each vertex index, while the last significant bits are for
   # invalid and deleted face
-  face: wp.array(dtype=int)
-  face_pr: wp.array(dtype=wp.vec3)
-  face_norm2: wp.array(dtype=float)
+  face: wp.array[int]
+  face_pr: wp.array[wp.vec3]
+  face_norm2: wp.array[float]
   nface: int
 
   # edges that make up the horizon when adding new vertices to polytope
-  horizon: wp.array(dtype=int)
+  horizon: wp.array[int]
   nhorizon: int
 
 
@@ -1335,7 +1335,7 @@ def _area4(a: wp.vec3, b: wp.vec3, c: wp.vec3, d: wp.vec3) -> float:
 
 
 @wp.func
-def _polygon_quad(polygon: wp.array(dtype=wp.vec3), npolygon: int) -> wp.vec4i:
+def _polygon_quad(polygon: wp.array[wp.vec3], npolygon: int) -> wp.vec4i:
   """Returns the indices of a quadrilateral of maximum area in a convex polygon (npolygon > 4)."""
   b = int(1)
   c = int(2)
@@ -1376,7 +1376,7 @@ def _polygon_quad(polygon: wp.array(dtype=wp.vec3), npolygon: int) -> wp.vec4i:
 # return number (1, 2 or 3) of dimensions of a simplex; reorder vertices if necessary
 @wp.func
 def _feature_dim(
-  face: wp.vec3i, vert_index: wp.array(dtype=int), vert: wp.array(dtype=wp.vec3), offset: int
+  face: wp.vec3i, vert_index: wp.array[int], vert: wp.array[wp.vec3], offset: int
 ) -> Tuple[int, wp.vec3i, wp.mat33]:
   v1i = vert_index[2 * face[0] + offset]
   v2i = vert_index[2 * face[1] + offset]
@@ -1401,9 +1401,7 @@ def _feature_dim(
 
 # find two normals that are facing each other within a tolerance, return 1 if found
 @wp.func
-def _aligned_faces(
-  vert1: wp.array(dtype=wp.vec3), len1: int, vert2: wp.array(dtype=wp.vec3), len2: int
-) -> Tuple[int, wp.vec2i]:
+def _aligned_faces(vert1: wp.array[wp.vec3], len1: int, vert2: wp.array[wp.vec3], len2: int) -> Tuple[int, wp.vec2i]:
   res = wp.vec2i()
   for i in range(len1):
     for j in range(len2):
@@ -1417,9 +1415,7 @@ def _aligned_faces(
 # find two normals that are perpendicular to each other within a tolerance
 # return 1 if found
 @wp.func
-def _aligned_face_edge(
-  edge: wp.array(dtype=wp.vec3), nedge: int, face: wp.array(dtype=wp.vec3), nface: int
-) -> Tuple[int, wp.vec2i]:
+def _aligned_face_edge(edge: wp.array[wp.vec3], nedge: int, face: wp.array[wp.vec3], nface: int) -> Tuple[int, wp.vec2i]:
   res = wp.vec2i()
   for i in range(nface):
     for j in range(nedge):
@@ -1432,9 +1428,7 @@ def _aligned_face_edge(
 
 # find up to n <= 2 common integers of two arrays, return n
 @wp.func
-def _intersect1(
-  a1: wp.array(dtype=int), a2: wp.array(dtype=int), start1: int, start2: int, len1: int, len2: int
-) -> Tuple[int, wp.vec2i]:
+def _intersect1(a1: wp.array[int], a2: wp.array[int], start1: int, start2: int, len1: int, len2: int) -> Tuple[int, wp.vec2i]:
   count = int(0)
   res = wp.vec2i()
   for i in range(start1, start1 + len1):
@@ -1448,7 +1442,7 @@ def _intersect1(
 
 
 @wp.func
-def _intersect2(a1: wp.vec2i, a2: wp.array(dtype=int), start2: int, len1: int, len2: int) -> Tuple[int, wp.vec2i]:
+def _intersect2(a1: wp.vec2i, a2: wp.array[int], start2: int, len1: int, len2: int) -> Tuple[int, wp.vec2i]:
   count = int(0)
   res = wp.vec2i()
   for i in range(len1):
@@ -1470,13 +1464,13 @@ def _mesh_normals(
   mat: wp.mat33,
   vertadr: int,
   polyadr: int,
-  polynormal: wp.array(dtype=wp.vec3),
-  polymapadr: wp.array(dtype=int),
-  polymapnum: wp.array(dtype=int),
-  polymap: wp.array(dtype=int),
+  polynormal: wp.array[wp.vec3],
+  polymapadr: wp.array[int],
+  polymapnum: wp.array[int],
+  polymap: wp.array[int],
   # Out:
-  normals_out: wp.array(dtype=wp.vec3),
-  indices_out: wp.array(dtype=int),
+  normals_out: wp.array[wp.vec3],
+  indices_out: wp.array[int],
 ) -> int:
   v1 = feature_index[0]
   v2 = feature_index[1]
@@ -1540,19 +1534,19 @@ def _mesh_edge_normals(
   pos: wp.vec3,
   vertadr: int,
   polyadr: int,
-  vert: wp.array(dtype=wp.vec3),
-  polyvertadr: wp.array(dtype=int),
-  polyvertnum: wp.array(dtype=int),
-  polyvert: wp.array(dtype=int),
-  polymapadr: wp.array(dtype=int),
-  polymapnum: wp.array(dtype=int),
-  polymap: wp.array(dtype=int),
+  vert: wp.array[wp.vec3],
+  polyvertadr: wp.array[int],
+  polyvertnum: wp.array[int],
+  polyvert: wp.array[int],
+  polymapadr: wp.array[int],
+  polymapnum: wp.array[int],
+  polymap: wp.array[int],
   v1: wp.vec3,
   v2: wp.vec3,
   v1i: int,
   # Out:
-  normals_out: wp.array(dtype=wp.vec3),
-  endverts_out: wp.array(dtype=wp.vec3),
+  normals_out: wp.array[wp.vec3],
+  endverts_out: wp.array[wp.vec3],
 ) -> int:
   # only one edge
   if dim == 2:
@@ -1586,8 +1580,8 @@ def _box_normals2(
   mat: wp.mat33,
   n: wp.vec3,
   # Out:
-  normal_out: wp.array(dtype=wp.vec3),
-  index_out: wp.array(dtype=int),
+  normal_out: wp.array[wp.vec3],
+  index_out: wp.array[int],
 ) -> int:
   # list of box face normals
   face_normals = mat63(1.0, 0.0, 0.0, -1.0, 0.0, 0.0, 0.0, 1.0, 0.0, 0.0, -1.0, 0.0, 0.0, 0.0, 1.0, 0.0, 0.0, -1.0)
@@ -1621,8 +1615,8 @@ def _box_normals(
   mat: wp.mat33,
   dir: wp.vec3,
   # Out:
-  normal_out: wp.array(dtype=wp.vec3),
-  index_out: wp.array(dtype=int),
+  normal_out: wp.array[wp.vec3],
+  index_out: wp.array[int],
 ) -> int:
   v1 = feature_index[0]
   v2 = feature_index[1]
@@ -1669,7 +1663,7 @@ def _box_normals(
     # c is 1 if edge is diagonal of a box face
     # c is 2 if edge is an external edge of box
     if c == 1 or c == 2:
-      return 2
+      return c
     return _box_normals2(mat, dir, normal_out, index_out)
 
   if feature_dim == 1:
@@ -1698,8 +1692,8 @@ def _box_edge_normals(
   v2: wp.vec3,
   v1i: int,
   # Out:
-  normal_out: wp.array(dtype=wp.vec3),
-  endvert_out: wp.array(dtype=wp.vec3),
+  normal_out: wp.array[wp.vec3],
+  endvert_out: wp.array[wp.vec3],
 ) -> int:
   if dim == 2:
     endvert_out[0] = v2
@@ -1726,7 +1720,7 @@ def _box_edge_normals(
 
 # recover face of a box from its index
 @wp.func
-def _box_face(mat: wp.mat33, pos: wp.vec3, size: wp.vec3, idx: int, face_out: wp.array(dtype=wp.vec3)) -> int:
+def _box_face(mat: wp.mat33, pos: wp.vec3, size: wp.vec3, idx: int, face_out: wp.array[wp.vec3]) -> int:
   # compute global coordinates of the box face and face normal
   if idx == 0:  # right
     face_out[0] = mat @ wp.vec3(size[0], size[1], size[2]) + pos
@@ -1775,13 +1769,13 @@ def _mesh_face(
   pos: wp.vec3,
   vertadr: int,
   polyadr: int,
-  vert: wp.array(dtype=wp.vec3),
-  polyvertadr: wp.array(dtype=int),
-  polyvertnum: wp.array(dtype=int),
-  polyvert: wp.array(dtype=int),
+  vert: wp.array[wp.vec3],
+  polyvertadr: wp.array[int],
+  polyvertnum: wp.array[int],
+  polyvert: wp.array[int],
   idx: int,
   # Out:
-  face_out: wp.array(dtype=wp.vec3),
+  face_out: wp.array[wp.vec3],
 ) -> int:
   adr = polyvertadr[polyadr + idx]
   j = int(0)
@@ -1821,17 +1815,17 @@ def _plane_intersect(pn: wp.vec3, pd: float, a: wp.vec3, b: wp.vec3) -> float:
 @wp.func
 def _polygon_clip(
   # In:
-  plane_normal: wp.array(dtype=wp.vec3),
-  plane_dist: wp.array(dtype=float),
-  face1: wp.array(dtype=wp.vec3),
+  plane_normal: wp.array[wp.vec3],
+  plane_dist: wp.array[float],
+  face1: wp.array[wp.vec3],
   nface1: int,
-  face2: wp.array(dtype=wp.vec3),
+  face2: wp.array[wp.vec3],
   nface2: int,
   n: wp.vec3,
   dir: wp.vec3,
   # Out:
-  polygon_out: wp.array(dtype=wp.vec3),
-  clipped_out: wp.array(dtype=wp.vec3),
+  polygon_out: wp.array[wp.vec3],
+  clipped_out: wp.array[wp.vec3],
 ) -> Tuple[int, mat43, mat43]:
   witness1 = mat43()
   witness2 = mat43()
@@ -1918,13 +1912,13 @@ def _polygon_clip(
 @wp.func
 def _set_edge(
   # In:
-  vert1: wp.array(dtype=wp.vec3),
-  vert2: wp.array(dtype=wp.vec3),
+  vert1: wp.array[wp.vec3],
+  vert2: wp.array[wp.vec3],
   start: int,
   end: int,
   offset: int,
   # Out:
-  face_out: wp.array(dtype=wp.vec3),
+  face_out: wp.array[wp.vec3],
 ) -> int:
   face_out[0] = vert1[2 * start + offset]
   face_out[1] = vert2[end]
@@ -1935,19 +1929,19 @@ def _set_edge(
 @wp.func
 def multicontact(
   # In:
-  polygon: wp.array(dtype=wp.vec3),
-  clipped: wp.array(dtype=wp.vec3),
-  plane_normal: wp.array(dtype=wp.vec3),
-  plane_dist: wp.array(dtype=float),
-  idx1: wp.array(dtype=int),
-  idx2: wp.array(dtype=int),
-  n1: wp.array(dtype=wp.vec3),
-  n2: wp.array(dtype=wp.vec3),
-  endvert: wp.array(dtype=wp.vec3),
-  face1: wp.array(dtype=wp.vec3),
-  face2: wp.array(dtype=wp.vec3),
-  epa_vert: wp.array(dtype=wp.vec3),
-  epa_vert_index: wp.array(dtype=int),
+  polygon: wp.array[wp.vec3],
+  clipped: wp.array[wp.vec3],
+  plane_normal: wp.array[wp.vec3],
+  plane_dist: wp.array[float],
+  idx1: wp.array[int],
+  idx2: wp.array[int],
+  n1: wp.array[wp.vec3],
+  n2: wp.array[wp.vec3],
+  endvert: wp.array[wp.vec3],
+  face1: wp.array[wp.vec3],
+  face2: wp.array[wp.vec3],
+  epa_vert: wp.array[wp.vec3],
+  epa_vert_index: wp.array[int],
   epa_face: int,
   x1: wp.vec3,
   x2: wp.vec3,
@@ -2215,12 +2209,12 @@ def ccd(
   geomtype2: int,
   x_1: wp.vec3,
   x_2: wp.vec3,
-  vert: wp.array(dtype=wp.vec3),
-  vert_index: wp.array(dtype=int),
-  face: wp.array(dtype=int),
-  face_pr: wp.array(dtype=wp.vec3),
-  face_norm2: wp.array(dtype=float),
-  horizon: wp.array(dtype=int),
+  vert: wp.array[wp.vec3],
+  vert_index: wp.array[int],
+  face: wp.array[int],
+  face_pr: wp.array[wp.vec3],
+  face_norm2: wp.array[float],
+  horizon: wp.array[int],
 ) -> Tuple[float, int, wp.vec3, wp.vec3, int]:
   """General convex collision detection via GJK/EPA."""
   full_margin1 = 0.0

mjx/mujoco/mjx/third_party/mujoco_warp/_src/collision_primitive.py

@@ -295,21 +295,21 @@ def plane_sphere_wrapper(
   geoms: wp.vec2i,
   pairid: wp.vec2i,
   # Data out:
-  contact_dist_out: wp.array(dtype=float),
-  contact_pos_out: wp.array(dtype=wp.vec3),
-  contact_frame_out: wp.array(dtype=wp.mat33),
-  contact_includemargin_out: wp.array(dtype=float),
-  contact_friction_out: wp.array(dtype=vec5),
-  contact_solref_out: wp.array(dtype=wp.vec2),
-  contact_solreffriction_out: wp.array(dtype=wp.vec2),
-  contact_solimp_out: wp.array(dtype=vec5),
-  contact_dim_out: wp.array(dtype=int),
-  contact_geom_out: wp.array(dtype=wp.vec2i),
-  contact_efc_address_out: wp.array2d(dtype=int),
-  contact_worldid_out: wp.array(dtype=int),
-  contact_type_out: wp.array(dtype=int),
-  contact_geomcollisionid_out: wp.array(dtype=int),
-  nacon_out: wp.array(dtype=int),
+  contact_dist_out: wp.array[float],
+  contact_pos_out: wp.array[wp.vec3],
+  contact_frame_out: wp.array[wp.mat33],
+  contact_includemargin_out: wp.array[float],
+  contact_friction_out: wp.array[vec5],
+  contact_solref_out: wp.array[wp.vec2],
+  contact_solreffriction_out: wp.array[wp.vec2],
+  contact_solimp_out: wp.array[vec5],
+  contact_dim_out: wp.array[int],
+  contact_geom_out: wp.array[wp.vec2i],
+  contact_efc_address_out: wp.array2d[int],
+  contact_worldid_out: wp.array[int],
+  contact_type_out: wp.array[int],
+  contact_geomcollisionid_out: wp.array[int],
+  nacon_out: wp.array[int],
 ):
   """Calculates contact between a sphere and a plane."""
   normal = plane.normal
@@ -367,21 +367,21 @@ def sphere_sphere_wrapper(
   geoms: wp.vec2i,
   pairid: wp.vec2i,
   # Data out:
-  contact_dist_out: wp.array(dtype=float),
-  contact_pos_out: wp.array(dtype=wp.vec3),
-  contact_frame_out: wp.array(dtype=wp.mat33),
-  contact_includemargin_out: wp.array(dtype=float),
-  contact_friction_out: wp.array(dtype=vec5),
-  contact_solref_out: wp.array(dtype=wp.vec2),
-  contact_solreffriction_out: wp.array(dtype=wp.vec2),
-  contact_solimp_out: wp.array(dtype=vec5),
-  contact_dim_out: wp.array(dtype=int),
-  contact_geom_out: wp.array(dtype=wp.vec2i),
-  contact_efc_address_out: wp.array2d(dtype=int),
-  contact_worldid_out: wp.array(dtype=int),
-  contact_type_out: wp.array(dtype=int),
-  contact_geomcollisionid_out: wp.array(dtype=int),
-  nacon_out: wp.array(dtype=int),
+  contact_dist_out: wp.array[float],
+  contact_pos_out: wp.array[wp.vec3],
+  contact_frame_out: wp.array[wp.mat33],
+  contact_includemargin_out: wp.array[float],
+  contact_friction_out: wp.array[vec5],
+  contact_solref_out: wp.array[wp.vec2],
+  contact_solreffriction_out: wp.array[wp.vec2],
+  contact_solimp_out: wp.array[vec5],
+  contact_dim_out: wp.array[int],
+  contact_geom_out: wp.array[wp.vec2i],
+  contact_efc_address_out: wp.array2d[int],
+  contact_worldid_out: wp.array[int],
+  contact_type_out: wp.array[int],
+  contact_geomcollisionid_out: wp.array[int],
+  nacon_out: wp.array[int],
 ):
   """Calculates contact between two spheres."""
   dist, pos, normal = sphere_sphere(sphere1.pos, sphere1.size[0], sphere2.pos, sphere2.size[0])
@@ -438,21 +438,21 @@ def sphere_capsule_wrapper(
   geoms: wp.vec2i,
   pairid: wp.vec2i,
   # Data out:
-  contact_dist_out: wp.array(dtype=float),
-  contact_pos_out: wp.array(dtype=wp.vec3),
-  contact_frame_out: wp.array(dtype=wp.mat33),
-  contact_includemargin_out: wp.array(dtype=float),
-  contact_friction_out: wp.array(dtype=vec5),
-  contact_solref_out: wp.array(dtype=wp.vec2),
-  contact_solreffriction_out: wp.array(dtype=wp.vec2),
-  contact_solimp_out: wp.array(dtype=vec5),
-  contact_dim_out: wp.array(dtype=int),
-  contact_geom_out: wp.array(dtype=wp.vec2i),
-  contact_efc_address_out: wp.array2d(dtype=int),
-  contact_worldid_out: wp.array(dtype=int),
-  contact_type_out: wp.array(dtype=int),
-  contact_geomcollisionid_out: wp.array(dtype=int),
-  nacon_out: wp.array(dtype=int),
+  contact_dist_out: wp.array[float],
+  contact_pos_out: wp.array[wp.vec3],
+  contact_frame_out: wp.array[wp.mat33],
+  contact_includemargin_out: wp.array[float],
+  contact_friction_out: wp.array[vec5],
+  contact_solref_out: wp.array[wp.vec2],
+  contact_solreffriction_out: wp.array[wp.vec2],
+  contact_solimp_out: wp.array[vec5],
+  contact_dim_out: wp.array[int],
+  contact_geom_out: wp.array[wp.vec2i],
+  contact_efc_address_out: wp.array2d[int],
+  contact_worldid_out: wp.array[int],
+  contact_type_out: wp.array[int],
+  contact_geomcollisionid_out: wp.array[int],
+  nacon_out: wp.array[int],
 ):
   """Calculates one contact between a sphere and a capsule."""
   # capsule axis
@@ -512,21 +512,21 @@ def capsule_capsule_wrapper(
   geoms: wp.vec2i,
   pairid: wp.vec2i,
   # Data out:
-  contact_dist_out: wp.array(dtype=float),
-  contact_pos_out: wp.array(dtype=wp.vec3),
-  contact_frame_out: wp.array(dtype=wp.mat33),
-  contact_includemargin_out: wp.array(dtype=float),
-  contact_friction_out: wp.array(dtype=vec5),
-  contact_solref_out: wp.array(dtype=wp.vec2),
-  contact_solreffriction_out: wp.array(dtype=wp.vec2),
-  contact_solimp_out: wp.array(dtype=vec5),
-  contact_dim_out: wp.array(dtype=int),
-  contact_geom_out: wp.array(dtype=wp.vec2i),
-  contact_efc_address_out: wp.array2d(dtype=int),
-  contact_worldid_out: wp.array(dtype=int),
-  contact_type_out: wp.array(dtype=int),
-  contact_geomcollisionid_out: wp.array(dtype=int),
-  nacon_out: wp.array(dtype=int),
+  contact_dist_out: wp.array[float],
+  contact_pos_out: wp.array[wp.vec3],
+  contact_frame_out: wp.array[wp.mat33],
+  contact_includemargin_out: wp.array[float],
+  contact_friction_out: wp.array[vec5],
+  contact_solref_out: wp.array[wp.vec2],
+  contact_solreffriction_out: wp.array[wp.vec2],
+  contact_solimp_out: wp.array[vec5],
+  contact_dim_out: wp.array[int],
+  contact_geom_out: wp.array[wp.vec2i],
+  contact_efc_address_out: wp.array2d[int],
+  contact_worldid_out: wp.array[int],
+  contact_type_out: wp.array[int],
+  contact_geomcollisionid_out: wp.array[int],
+  nacon_out: wp.array[int],
 ):
   """Calculates contacts between two capsules."""
   # capsule axes
@@ -598,21 +598,21 @@ def plane_capsule_wrapper(
   geoms: wp.vec2i,
   pairid: wp.vec2i,
   # Data out:
-  contact_dist_out: wp.array(dtype=float),
-  contact_pos_out: wp.array(dtype=wp.vec3),
-  contact_frame_out: wp.array(dtype=wp.mat33),
-  contact_includemargin_out: wp.array(dtype=float),
-  contact_friction_out: wp.array(dtype=vec5),
-  contact_solref_out: wp.array(dtype=wp.vec2),
-  contact_solreffriction_out: wp.array(dtype=wp.vec2),
-  contact_solimp_out: wp.array(dtype=vec5),
-  contact_dim_out: wp.array(dtype=int),
-  contact_geom_out: wp.array(dtype=wp.vec2i),
-  contact_efc_address_out: wp.array2d(dtype=int),
-  contact_worldid_out: wp.array(dtype=int),
-  contact_type_out: wp.array(dtype=int),
-  contact_geomcollisionid_out: wp.array(dtype=int),
-  nacon_out: wp.array(dtype=int),
+  contact_dist_out: wp.array[float],
+  contact_pos_out: wp.array[wp.vec3],
+  contact_frame_out: wp.array[wp.mat33],
+  contact_includemargin_out: wp.array[float],
+  contact_friction_out: wp.array[vec5],
+  contact_solref_out: wp.array[wp.vec2],
+  contact_solreffriction_out: wp.array[wp.vec2],
+  contact_solimp_out: wp.array[vec5],
+  contact_dim_out: wp.array[int],
+  contact_geom_out: wp.array[wp.vec2i],
+  contact_efc_address_out: wp.array2d[int],
+  contact_worldid_out: wp.array[int],
+  contact_type_out: wp.array[int],
+  contact_geomcollisionid_out: wp.array[int],
+  nacon_out: wp.array[int],
 ):
   """Calculates contacts between a capsule and a plane."""
   # capsule axis
@@ -680,21 +680,21 @@ def plane_ellipsoid_wrapper(
   geoms: wp.vec2i,
   pairid: wp.vec2i,
   # Data out:
-  contact_dist_out: wp.array(dtype=float),
-  contact_pos_out: wp.array(dtype=wp.vec3),
-  contact_frame_out: wp.array(dtype=wp.mat33),
-  contact_includemargin_out: wp.array(dtype=float),
-  contact_friction_out: wp.array(dtype=vec5),
-  contact_solref_out: wp.array(dtype=wp.vec2),
-  contact_solreffriction_out: wp.array(dtype=wp.vec2),
-  contact_solimp_out: wp.array(dtype=vec5),
-  contact_dim_out: wp.array(dtype=int),
-  contact_geom_out: wp.array(dtype=wp.vec2i),
-  contact_efc_address_out: wp.array2d(dtype=int),
-  contact_worldid_out: wp.array(dtype=int),
-  contact_type_out: wp.array(dtype=int),
-  contact_geomcollisionid_out: wp.array(dtype=int),
-  nacon_out: wp.array(dtype=int),
+  contact_dist_out: wp.array[float],
+  contact_pos_out: wp.array[wp.vec3],
+  contact_frame_out: wp.array[wp.mat33],
+  contact_includemargin_out: wp.array[float],
+  contact_friction_out: wp.array[vec5],
+  contact_solref_out: wp.array[wp.vec2],
+  contact_solreffriction_out: wp.array[wp.vec2],
+  contact_solimp_out: wp.array[vec5],
+  contact_dim_out: wp.array[int],
+  contact_geom_out: wp.array[wp.vec2i],
+  contact_efc_address_out: wp.array2d[int],
+  contact_worldid_out: wp.array[int],
+  contact_type_out: wp.array[int],
+  contact_geomcollisionid_out: wp.array[int],
+  nacon_out: wp.array[int],
 ):
   """Calculates contacts between an ellipsoid and a plane."""
   dist, pos, normal = plane_ellipsoid(plane.normal, plane.pos, ellipsoid.pos, ellipsoid.rot, ellipsoid.size)
@@ -751,21 +751,21 @@ def plane_box_wrapper(
   geoms: wp.vec2i,
   pairid: wp.vec2i,
   # Data out:
-  contact_dist_out: wp.array(dtype=float),
-  contact_pos_out: wp.array(dtype=wp.vec3),
-  contact_frame_out: wp.array(dtype=wp.mat33),
-  contact_includemargin_out: wp.array(dtype=float),
-  contact_friction_out: wp.array(dtype=vec5),
-  contact_solref_out: wp.array(dtype=wp.vec2),
-  contact_solreffriction_out: wp.array(dtype=wp.vec2),
-  contact_solimp_out: wp.array(dtype=vec5),
-  contact_dim_out: wp.array(dtype=int),
-  contact_geom_out: wp.array(dtype=wp.vec2i),
-  contact_efc_address_out: wp.array2d(dtype=int),
-  contact_worldid_out: wp.array(dtype=int),
-  contact_type_out: wp.array(dtype=int),
-  contact_geomcollisionid_out: wp.array(dtype=int),
-  nacon_out: wp.array(dtype=int),
+  contact_dist_out: wp.array[float],
+  contact_pos_out: wp.array[wp.vec3],
+  contact_frame_out: wp.array[wp.mat33],
+  contact_includemargin_out: wp.array[float],
+  contact_friction_out: wp.array[vec5],
+  contact_solref_out: wp.array[wp.vec2],
+  contact_solreffriction_out: wp.array[wp.vec2],
+  contact_solimp_out: wp.array[vec5],
+  contact_dim_out: wp.array[int],
+  contact_geom_out: wp.array[wp.vec2i],
+  contact_efc_address_out: wp.array2d[int],
+  contact_worldid_out: wp.array[int],
+  contact_type_out: wp.array[int],
+  contact_geomcollisionid_out: wp.array[int],
+  nacon_out: wp.array[int],
 ):
   """Calculates contacts between a box and a plane."""
   dist, pos, normal = plane_box(plane.normal, plane.pos, box.pos, box.rot, box.size)
@@ -824,21 +824,21 @@ def plane_convex_wrapper(
   geoms: wp.vec2i,
   pairid: wp.vec2i,
   # Data out:
-  contact_dist_out: wp.array(dtype=float),
-  contact_pos_out: wp.array(dtype=wp.vec3),
-  contact_frame_out: wp.array(dtype=wp.mat33),
-  contact_includemargin_out: wp.array(dtype=float),
-  contact_friction_out: wp.array(dtype=vec5),
-  contact_solref_out: wp.array(dtype=wp.vec2),
-  contact_solreffriction_out: wp.array(dtype=wp.vec2),
-  contact_solimp_out: wp.array(dtype=vec5),
-  contact_dim_out: wp.array(dtype=int),
-  contact_geom_out: wp.array(dtype=wp.vec2i),
-  contact_efc_address_out: wp.array2d(dtype=int),
-  contact_worldid_out: wp.array(dtype=int),
-  contact_type_out: wp.array(dtype=int),
-  contact_geomcollisionid_out: wp.array(dtype=int),
-  nacon_out: wp.array(dtype=int),
+  contact_dist_out: wp.array[float],
+  contact_pos_out: wp.array[wp.vec3],
+  contact_frame_out: wp.array[wp.mat33],
+  contact_includemargin_out: wp.array[float],
+  contact_friction_out: wp.array[vec5],
+  contact_solref_out: wp.array[wp.vec2],
+  contact_solreffriction_out: wp.array[wp.vec2],
+  contact_solimp_out: wp.array[vec5],
+  contact_dim_out: wp.array[int],
+  contact_geom_out: wp.array[wp.vec2i],
+  contact_efc_address_out: wp.array2d[int],
+  contact_worldid_out: wp.array[int],
+  contact_type_out: wp.array[int],
+  contact_geomcollisionid_out: wp.array[int],
+  nacon_out: wp.array[int],
 ):
   """Calculates contacts between a plane and a convex object."""
   dist, pos, normal = plane_convex(plane.normal, plane.pos, convex)
@@ -897,21 +897,21 @@ def sphere_cylinder_wrapper(
   geoms: wp.vec2i,
   pairid: wp.vec2i,
   # Data out:
-  contact_dist_out: wp.array(dtype=float),
-  contact_pos_out: wp.array(dtype=wp.vec3),
-  contact_frame_out: wp.array(dtype=wp.mat33),
-  contact_includemargin_out: wp.array(dtype=float),
-  contact_friction_out: wp.array(dtype=vec5),
-  contact_solref_out: wp.array(dtype=wp.vec2),
-  contact_solreffriction_out: wp.array(dtype=wp.vec2),
-  contact_solimp_out: wp.array(dtype=vec5),
-  contact_dim_out: wp.array(dtype=int),
-  contact_geom_out: wp.array(dtype=wp.vec2i),
-  contact_efc_address_out: wp.array2d(dtype=int),
-  contact_worldid_out: wp.array(dtype=int),
-  contact_type_out: wp.array(dtype=int),
-  contact_geomcollisionid_out: wp.array(dtype=int),
-  nacon_out: wp.array(dtype=int),
+  contact_dist_out: wp.array[float],
+  contact_pos_out: wp.array[wp.vec3],
+  contact_frame_out: wp.array[wp.mat33],
+  contact_includemargin_out: wp.array[float],
+  contact_friction_out: wp.array[vec5],
+  contact_solref_out: wp.array[wp.vec2],
+  contact_solreffriction_out: wp.array[wp.vec2],
+  contact_solimp_out: wp.array[vec5],
+  contact_dim_out: wp.array[int],
+  contact_geom_out: wp.array[wp.vec2i],
+  contact_efc_address_out: wp.array2d[int],
+  contact_worldid_out: wp.array[int],
+  contact_type_out: wp.array[int],
+  contact_geomcollisionid_out: wp.array[int],
+  nacon_out: wp.array[int],
 ):
   """Calculates contacts between a sphere and a cylinder."""
   # cylinder axis
@@ -978,21 +978,21 @@ def plane_cylinder_wrapper(
   geoms: wp.vec2i,
   pairid: wp.vec2i,
   # Data out:
-  contact_dist_out: wp.array(dtype=float),
-  contact_pos_out: wp.array(dtype=wp.vec3),
-  contact_frame_out: wp.array(dtype=wp.mat33),
-  contact_includemargin_out: wp.array(dtype=float),
-  contact_friction_out: wp.array(dtype=vec5),
-  contact_solref_out: wp.array(dtype=wp.vec2),
-  contact_solreffriction_out: wp.array(dtype=wp.vec2),
-  contact_solimp_out: wp.array(dtype=vec5),
-  contact_dim_out: wp.array(dtype=int),
-  contact_geom_out: wp.array(dtype=wp.vec2i),
-  contact_efc_address_out: wp.array2d(dtype=int),
-  contact_worldid_out: wp.array(dtype=int),
-  contact_type_out: wp.array(dtype=int),
-  contact_geomcollisionid_out: wp.array(dtype=int),
-  nacon_out: wp.array(dtype=int),
+  contact_dist_out: wp.array[float],
+  contact_pos_out: wp.array[wp.vec3],
+  contact_frame_out: wp.array[wp.mat33],
+  contact_includemargin_out: wp.array[float],
+  contact_friction_out: wp.array[vec5],
+  contact_solref_out: wp.array[wp.vec2],
+  contact_solreffriction_out: wp.array[wp.vec2],
+  contact_solimp_out: wp.array[vec5],
+  contact_dim_out: wp.array[int],
+  contact_geom_out: wp.array[wp.vec2i],
+  contact_efc_address_out: wp.array2d[int],
+  contact_worldid_out: wp.array[int],
+  contact_type_out: wp.array[int],
+  contact_geomcollisionid_out: wp.array[int],
+  nacon_out: wp.array[int],
 ):
   """Calculates contacts between a cylinder and a plane."""
   # cylinder axis
@@ -1061,21 +1061,21 @@ def sphere_box_wrapper(
   geoms: wp.vec2i,
   pairid: wp.vec2i,
   # Data out:
-  contact_dist_out: wp.array(dtype=float),
-  contact_pos_out: wp.array(dtype=wp.vec3),
-  contact_frame_out: wp.array(dtype=wp.mat33),
-  contact_includemargin_out: wp.array(dtype=float),
-  contact_friction_out: wp.array(dtype=vec5),
-  contact_solref_out: wp.array(dtype=wp.vec2),
-  contact_solreffriction_out: wp.array(dtype=wp.vec2),
-  contact_solimp_out: wp.array(dtype=vec5),
-  contact_dim_out: wp.array(dtype=int),
-  contact_geom_out: wp.array(dtype=wp.vec2i),
-  contact_efc_address_out: wp.array2d(dtype=int),
-  contact_worldid_out: wp.array(dtype=int),
-  contact_type_out: wp.array(dtype=int),
-  contact_geomcollisionid_out: wp.array(dtype=int),
-  nacon_out: wp.array(dtype=int),
+  contact_dist_out: wp.array[float],
+  contact_pos_out: wp.array[wp.vec3],
+  contact_frame_out: wp.array[wp.mat33],
+  contact_includemargin_out: wp.array[float],
+  contact_friction_out: wp.array[vec5],
+  contact_solref_out: wp.array[wp.vec2],
+  contact_solreffriction_out: wp.array[wp.vec2],
+  contact_solimp_out: wp.array[vec5],
+  contact_dim_out: wp.array[int],
+  contact_geom_out: wp.array[wp.vec2i],
+  contact_efc_address_out: wp.array2d[int],
+  contact_worldid_out: wp.array[int],
+  contact_type_out: wp.array[int],
+  contact_geomcollisionid_out: wp.array[int],
+  nacon_out: wp.array[int],
 ):
   dist, pos, normal = sphere_box(sphere.pos, sphere.size[0], box.pos, box.rot, box.size)
 
@@ -1131,21 +1131,21 @@ def capsule_box_wrapper(
   geoms: wp.vec2i,
   pairid: wp.vec2i,
   # Data out:
-  contact_dist_out: wp.array(dtype=float),
-  contact_pos_out: wp.array(dtype=wp.vec3),
-  contact_frame_out: wp.array(dtype=wp.mat33),
-  contact_includemargin_out: wp.array(dtype=float),
-  contact_friction_out: wp.array(dtype=vec5),
-  contact_solref_out: wp.array(dtype=wp.vec2),
-  contact_solreffriction_out: wp.array(dtype=wp.vec2),
-  contact_solimp_out: wp.array(dtype=vec5),
-  contact_dim_out: wp.array(dtype=int),
-  contact_geom_out: wp.array(dtype=wp.vec2i),
-  contact_efc_address_out: wp.array2d(dtype=int),
-  contact_worldid_out: wp.array(dtype=int),
-  contact_type_out: wp.array(dtype=int),
-  contact_geomcollisionid_out: wp.array(dtype=int),
-  nacon_out: wp.array(dtype=int),
+  contact_dist_out: wp.array[float],
+  contact_pos_out: wp.array[wp.vec3],
+  contact_frame_out: wp.array[wp.mat33],
+  contact_includemargin_out: wp.array[float],
+  contact_friction_out: wp.array[vec5],
+  contact_solref_out: wp.array[wp.vec2],
+  contact_solreffriction_out: wp.array[wp.vec2],
+  contact_solimp_out: wp.array[vec5],
+  contact_dim_out: wp.array[int],
+  contact_geom_out: wp.array[wp.vec2i],
+  contact_efc_address_out: wp.array2d[int],
+  contact_worldid_out: wp.array[int],
+  contact_type_out: wp.array[int],
+  contact_geomcollisionid_out: wp.array[int],
+  nacon_out: wp.array[int],
 ):
   """Calculates contacts between a capsule and a box."""
   # Extract capsule axis
@@ -1216,21 +1216,21 @@ def box_box_wrapper(
   geoms: wp.vec2i,
   pairid: wp.vec2i,
   # Data out:
-  contact_dist_out: wp.array(dtype=float),
-  contact_pos_out: wp.array(dtype=wp.vec3),
-  contact_frame_out: wp.array(dtype=wp.mat33),
-  contact_includemargin_out: wp.array(dtype=float),
-  contact_friction_out: wp.array(dtype=vec5),
-  contact_solref_out: wp.array(dtype=wp.vec2),
-  contact_solreffriction_out: wp.array(dtype=wp.vec2),
-  contact_solimp_out: wp.array(dtype=vec5),
-  contact_dim_out: wp.array(dtype=int),
-  contact_geom_out: wp.array(dtype=wp.vec2i),
-  contact_efc_address_out: wp.array2d(dtype=int),
-  contact_worldid_out: wp.array(dtype=int),
-  contact_type_out: wp.array(dtype=int),
-  contact_geomcollisionid_out: wp.array(dtype=int),
-  nacon_out: wp.array(dtype=int),
+  contact_dist_out: wp.array[float],
+  contact_pos_out: wp.array[wp.vec3],
+  contact_frame_out: wp.array[wp.mat33],
+  contact_includemargin_out: wp.array[float],
+  contact_friction_out: wp.array[vec5],
+  contact_solref_out: wp.array[wp.vec2],
+  contact_solreffriction_out: wp.array[wp.vec2],
+  contact_solimp_out: wp.array[vec5],
+  contact_dim_out: wp.array[int],
+  contact_geom_out: wp.array[wp.vec2i],
+  contact_efc_address_out: wp.array2d[int],
+  contact_worldid_out: wp.array[int],
+  contact_type_out: wp.array[int],
+  contact_geomcollisionid_out: wp.array[int],
+  nacon_out: wp.array[int],
 ):
   """Calculates contacts between two boxes."""
   # Call the core function to get contact geometry
@@ -1302,63 +1302,63 @@ def _primitive_narrowphase(primitive_collisions_types, primitive_collisions_func
   @wp.kernel(module="unique", enable_backward=False)
   def primitive_narrowphase(
     # Model:
-    geom_type: wp.array(dtype=int),
-    geom_condim: wp.array(dtype=int),
-    geom_dataid: wp.array(dtype=int),
-    geom_priority: wp.array(dtype=int),
-    geom_solmix: wp.array2d(dtype=float),
-    geom_solref: wp.array2d(dtype=wp.vec2),
-    geom_solimp: wp.array2d(dtype=vec5),
-    geom_size: wp.array2d(dtype=wp.vec3),
-    geom_friction: wp.array2d(dtype=wp.vec3),
-    geom_margin: wp.array2d(dtype=float),
-    geom_gap: wp.array2d(dtype=float),
-    mesh_vertadr: wp.array(dtype=int),
-    mesh_vertnum: wp.array(dtype=int),
-    mesh_graphadr: wp.array(dtype=int),
-    mesh_vert: wp.array(dtype=wp.vec3),
-    mesh_graph: wp.array(dtype=int),
-    mesh_polynum: wp.array(dtype=int),
-    mesh_polyadr: wp.array(dtype=int),
-    mesh_polynormal: wp.array(dtype=wp.vec3),
-    mesh_polyvertadr: wp.array(dtype=int),
-    mesh_polyvertnum: wp.array(dtype=int),
-    mesh_polyvert: wp.array(dtype=int),
-    mesh_polymapadr: wp.array(dtype=int),
-    mesh_polymapnum: wp.array(dtype=int),
-    mesh_polymap: wp.array(dtype=int),
-    pair_dim: wp.array(dtype=int),
-    pair_solref: wp.array2d(dtype=wp.vec2),
-    pair_solreffriction: wp.array2d(dtype=wp.vec2),
-    pair_solimp: wp.array2d(dtype=vec5),
-    pair_margin: wp.array2d(dtype=float),
-    pair_gap: wp.array2d(dtype=float),
-    pair_friction: wp.array2d(dtype=vec5),
+    geom_type: wp.array[int],
+    geom_condim: wp.array[int],
+    geom_dataid: wp.array2d[int],
+    geom_priority: wp.array[int],
+    geom_solmix: wp.array2d[float],
+    geom_solref: wp.array2d[wp.vec2],
+    geom_solimp: wp.array2d[vec5],
+    geom_size: wp.array2d[wp.vec3],
+    geom_friction: wp.array2d[wp.vec3],
+    geom_margin: wp.array2d[float],
+    geom_gap: wp.array2d[float],
+    mesh_vertadr: wp.array[int],
+    mesh_vertnum: wp.array[int],
+    mesh_graphadr: wp.array[int],
+    mesh_vert: wp.array[wp.vec3],
+    mesh_graph: wp.array[int],
+    mesh_polynum: wp.array[int],
+    mesh_polyadr: wp.array[int],
+    mesh_polynormal: wp.array[wp.vec3],
+    mesh_polyvertadr: wp.array[int],
+    mesh_polyvertnum: wp.array[int],
+    mesh_polyvert: wp.array[int],
+    mesh_polymapadr: wp.array[int],
+    mesh_polymapnum: wp.array[int],
+    mesh_polymap: wp.array[int],
+    pair_dim: wp.array[int],
+    pair_solref: wp.array2d[wp.vec2],
+    pair_solreffriction: wp.array2d[wp.vec2],
+    pair_solimp: wp.array2d[vec5],
+    pair_margin: wp.array2d[float],
+    pair_gap: wp.array2d[float],
+    pair_friction: wp.array2d[vec5],
     # Data in:
-    geom_xpos_in: wp.array2d(dtype=wp.vec3),
-    geom_xmat_in: wp.array2d(dtype=wp.mat33),
+    geom_xpos_in: wp.array2d[wp.vec3],
+    geom_xmat_in: wp.array2d[wp.mat33],
     naconmax_in: int,
-    ncollision_in: wp.array(dtype=int),
+    ncollision_in: wp.array[int],
     # In:
-    collision_pair_in: wp.array(dtype=wp.vec2i),
-    collision_pairid_in: wp.array(dtype=wp.vec2i),
-    collision_worldid_in: wp.array(dtype=int),
+    collision_pair_in: wp.array[wp.vec2i],
+    collision_pairid_in: wp.array[wp.vec2i],
+    collision_worldid_in: wp.array[int],
     # Data out:
-    contact_dist_out: wp.array(dtype=float),
-    contact_pos_out: wp.array(dtype=wp.vec3),
-    contact_frame_out: wp.array(dtype=wp.mat33),
-    contact_includemargin_out: wp.array(dtype=float),
-    contact_friction_out: wp.array(dtype=vec5),
-    contact_solref_out: wp.array(dtype=wp.vec2),
-    contact_solreffriction_out: wp.array(dtype=wp.vec2),
-    contact_solimp_out: wp.array(dtype=vec5),
-    contact_dim_out: wp.array(dtype=int),
-    contact_geom_out: wp.array(dtype=wp.vec2i),
-    contact_efc_address_out: wp.array2d(dtype=int),
-    contact_worldid_out: wp.array(dtype=int),
-    contact_type_out: wp.array(dtype=int),
-    contact_geomcollisionid_out: wp.array(dtype=int),
-    nacon_out: wp.array(dtype=int),
+    contact_dist_out: wp.array[float],
+    contact_pos_out: wp.array[wp.vec3],
+    contact_frame_out: wp.array[wp.mat33],
+    contact_includemargin_out: wp.array[float],
+    contact_friction_out: wp.array[vec5],
+    contact_solref_out: wp.array[wp.vec2],
+    contact_solreffriction_out: wp.array[wp.vec2],
+    contact_solimp_out: wp.array[vec5],
+    contact_dim_out: wp.array[int],
+    contact_geom_out: wp.array[wp.vec2i],
+    contact_efc_address_out: wp.array2d[int],
+    contact_worldid_out: wp.array[int],
+    contact_type_out: wp.array[int],
+    contact_geomcollisionid_out: wp.array[int],
+    nacon_out: wp.array[int],
   ):
     tid = wp.tid()
 

mjx/mujoco/mjx/third_party/mujoco_warp/_src/collision_sdf.py

@@ -54,9 +54,9 @@ class AABB:
 class VolumeData:
   center: wp.vec3
   half_size: wp.vec3
-  oct_aabb: wp.array2d(dtype=wp.vec3)
-  oct_child: wp.array(dtype=vec8i)
-  oct_coeff: wp.array(dtype=vec8)
+  oct_aabb: wp.array2d[wp.vec3]
+  oct_child: wp.array[vec8i]
+  oct_coeff: wp.array[vec8]
   root: int = 0
   valid: bool = False
 
@@ -64,10 +64,10 @@ class VolumeData:
 @wp.struct
 class MeshData:
   nmeshface: int
-  mesh_vertadr: wp.array(dtype=int)
-  mesh_vert: wp.array(dtype=wp.vec3)
-  mesh_faceadr: wp.array(dtype=int)
-  mesh_face: wp.array(dtype=wp.vec3i)
+  mesh_vertadr: wp.array[int]
+  mesh_vert: wp.array[wp.vec3]
+  mesh_faceadr: wp.array[int]
+  mesh_face: wp.array[wp.vec3i]
   data_id: int
   pos: wp.vec3
   mat: wp.mat33
@@ -80,12 +80,12 @@ class MeshData:
 @wp.func
 def get_sdf_params(
   # Model:
-  oct_child: wp.array(dtype=vec8i),
-  oct_aabb: wp.array2d(dtype=wp.vec3),
-  oct_coeff: wp.array(dtype=vec8),
-  mesh_octadr: wp.array(dtype=int),
-  plugin: wp.array(dtype=int),
-  plugin_attr: wp.array(dtype=vec_pluginattr),
+  oct_child: wp.array[vec8i],
+  oct_aabb: wp.array2d[wp.vec3],
+  oct_coeff: wp.array[vec8],
+  mesh_octadr: wp.array[int],
+  plugin: wp.array[int],
+  plugin_attr: wp.array[vec_pluginattr],
   # In:
   g_type: int,
   g_size: wp.vec3,
@@ -252,7 +252,7 @@ def user_sdf_grad(p: wp.vec3, attr: vec_pluginattr, sdf_type: int) -> wp.vec3:
 
 @wp.func
 def find_oct(
-  oct_child: wp.array(dtype=vec8i), oct_aabb: wp.array2d(dtype=wp.vec3), p: wp.vec3, grad: bool, root: int
+  oct_child: wp.array[vec8i], oct_aabb: wp.array2d[wp.vec3], p: wp.vec3, grad: bool, root: int
 ) -> Tuple[int, Tuple[vec8, vec8, vec8]]:
   stack = root
   niter = int(100)
@@ -665,75 +665,75 @@ def gradient_descent(
 def _sdf_narrowphase(
   # Model:
   nmeshface: int,
-  oct_child: wp.array(dtype=vec8i),
-  oct_aabb: wp.array2d(dtype=wp.vec3),
-  oct_coeff: wp.array(dtype=vec8),
-  geom_type: wp.array(dtype=int),
-  geom_condim: wp.array(dtype=int),
-  geom_dataid: wp.array(dtype=int),
-  geom_priority: wp.array(dtype=int),
-  geom_solmix: wp.array2d(dtype=float),
-  geom_solref: wp.array2d(dtype=wp.vec2),
-  geom_solimp: wp.array2d(dtype=vec5),
-  geom_size: wp.array2d(dtype=wp.vec3),
-  geom_aabb: wp.array3d(dtype=wp.vec3),
-  geom_friction: wp.array2d(dtype=wp.vec3),
-  geom_margin: wp.array2d(dtype=float),
-  geom_gap: wp.array2d(dtype=float),
-  mesh_vertadr: wp.array(dtype=int),
-  mesh_vertnum: wp.array(dtype=int),
-  mesh_faceadr: wp.array(dtype=int),
-  mesh_octadr: wp.array(dtype=int),
-  mesh_graphadr: wp.array(dtype=int),
-  mesh_vert: wp.array(dtype=wp.vec3),
-  mesh_face: wp.array(dtype=wp.vec3i),
-  mesh_graph: wp.array(dtype=int),
-  mesh_polynum: wp.array(dtype=int),
-  mesh_polyadr: wp.array(dtype=int),
-  mesh_polynormal: wp.array(dtype=wp.vec3),
-  mesh_polyvertadr: wp.array(dtype=int),
-  mesh_polyvertnum: wp.array(dtype=int),
-  mesh_polyvert: wp.array(dtype=int),
-  mesh_polymapadr: wp.array(dtype=int),
-  mesh_polymapnum: wp.array(dtype=int),
-  mesh_polymap: wp.array(dtype=int),
-  pair_dim: wp.array(dtype=int),
-  pair_solref: wp.array2d(dtype=wp.vec2),
-  pair_solreffriction: wp.array2d(dtype=wp.vec2),
-  pair_solimp: wp.array2d(dtype=vec5),
-  pair_margin: wp.array2d(dtype=float),
-  pair_gap: wp.array2d(dtype=float),
-  pair_friction: wp.array2d(dtype=vec5),
-  plugin: wp.array(dtype=int),
-  plugin_attr: wp.array(dtype=vec_pluginattr),
-  geom_plugin_index: wp.array(dtype=int),
+  oct_child: wp.array[vec8i],
+  oct_aabb: wp.array2d[wp.vec3],
+  oct_coeff: wp.array[vec8],
+  geom_type: wp.array[int],
+  geom_condim: wp.array[int],
+  geom_dataid: wp.array2d[int],
+  geom_priority: wp.array[int],
+  geom_solmix: wp.array2d[float],
+  geom_solref: wp.array2d[wp.vec2],
+  geom_solimp: wp.array2d[vec5],
+  geom_size: wp.array2d[wp.vec3],
+  geom_aabb: wp.array3d[wp.vec3],
+  geom_friction: wp.array2d[wp.vec3],
+  geom_margin: wp.array2d[float],
+  geom_gap: wp.array2d[float],
+  mesh_vertadr: wp.array[int],
+  mesh_vertnum: wp.array[int],
+  mesh_faceadr: wp.array[int],
+  mesh_octadr: wp.array[int],
+  mesh_graphadr: wp.array[int],
+  mesh_vert: wp.array[wp.vec3],
+  mesh_face: wp.array[wp.vec3i],
+  mesh_graph: wp.array[int],
+  mesh_polynum: wp.array[int],
+  mesh_polyadr: wp.array[int],
+  mesh_polynormal: wp.array[wp.vec3],
+  mesh_polyvertadr: wp.array[int],
+  mesh_polyvertnum: wp.array[int],
+  mesh_polyvert: wp.array[int],
+  mesh_polymapadr: wp.array[int],
+  mesh_polymapnum: wp.array[int],
+  mesh_polymap: wp.array[int],
+  pair_dim: wp.array[int],
+  pair_solref: wp.array2d[wp.vec2],
+  pair_solreffriction: wp.array2d[wp.vec2],
+  pair_solimp: wp.array2d[vec5],
+  pair_margin: wp.array2d[float],
+  pair_gap: wp.array2d[float],
+  pair_friction: wp.array2d[vec5],
+  plugin: wp.array[int],
+  plugin_attr: wp.array[vec_pluginattr],
+  geom_plugin_index: wp.array[int],
   # Data in:
-  geom_xpos_in: wp.array2d(dtype=wp.vec3),
-  geom_xmat_in: wp.array2d(dtype=wp.mat33),
+  geom_xpos_in: wp.array2d[wp.vec3],
+  geom_xmat_in: wp.array2d[wp.mat33],
   naconmax_in: int,
-  ncollision_in: wp.array(dtype=int),
+  ncollision_in: wp.array[int],
   # In:
-  collision_pair_in: wp.array(dtype=wp.vec2i),
-  collision_pairid_in: wp.array(dtype=wp.vec2i),
-  collision_worldid_in: wp.array(dtype=int),
+  collision_pair_in: wp.array[wp.vec2i],
+  collision_pairid_in: wp.array[wp.vec2i],
+  collision_worldid_in: wp.array[int],
   sdf_initpoints: int,
   sdf_iterations: int,
   # Data out:
-  contact_dist_out: wp.array(dtype=float),
-  contact_pos_out: wp.array(dtype=wp.vec3),
-  contact_frame_out: wp.array(dtype=wp.mat33),
-  contact_includemargin_out: wp.array(dtype=float),
-  contact_friction_out: wp.array(dtype=vec5),
-  contact_solref_out: wp.array(dtype=wp.vec2),
-  contact_solreffriction_out: wp.array(dtype=wp.vec2),
-  contact_solimp_out: wp.array(dtype=vec5),
-  contact_dim_out: wp.array(dtype=int),
-  contact_geom_out: wp.array(dtype=wp.vec2i),
-  contact_efc_address_out: wp.array2d(dtype=int),
-  contact_worldid_out: wp.array(dtype=int),
-  contact_type_out: wp.array(dtype=int),
-  contact_geomcollisionid_out: wp.array(dtype=int),
-  nacon_out: wp.array(dtype=int),
+  contact_dist_out: wp.array[float],
+  contact_pos_out: wp.array[wp.vec3],
+  contact_frame_out: wp.array[wp.mat33],
+  contact_includemargin_out: wp.array[float],
+  contact_friction_out: wp.array[vec5],
+  contact_solref_out: wp.array[wp.vec2],
+  contact_solreffriction_out: wp.array[wp.vec2],
+  contact_solimp_out: wp.array[vec5],
+  contact_dim_out: wp.array[int],
+  contact_geom_out: wp.array[wp.vec2i],
+  contact_efc_address_out: wp.array2d[int],
+  contact_worldid_out: wp.array[int],
+  contact_type_out: wp.array[int],
+  contact_geomcollisionid_out: wp.array[int],
+  nacon_out: wp.array[int],
 ):
   i, contact_tid = wp.tid()
   if i >= sdf_initpoints:
@@ -817,12 +817,32 @@ def _sdf_narrowphase(
   pos1 = geom1.pos
   rot1 = geom1.rot
 
+  dataid_setid = worldid % geom_dataid.shape[0]
+
   attr1, g1_plugin_id, volume_data1, mesh_data1 = get_sdf_params(
-    oct_child, oct_aabb, oct_coeff, mesh_octadr, plugin, plugin_attr, type1, geom1.size, g1_plugin, geom_dataid[g1]
+    oct_child,
+    oct_aabb,
+    oct_coeff,
+    mesh_octadr,
+    plugin,
+    plugin_attr,
+    type1,
+    geom1.size,
+    g1_plugin,
+    geom_dataid[dataid_setid, g1],
   )
 
   attr2, g2_plugin_id, volume_data2, mesh_data2 = get_sdf_params(
-    oct_child, oct_aabb, oct_coeff, mesh_octadr, plugin, plugin_attr, type2, geom2.size, g2_plugin, geom_dataid[g2]
+    oct_child,
+    oct_aabb,
+    oct_coeff,
+    mesh_octadr,
+    plugin,
+    plugin_attr,
+    type2,
+    geom2.size,
+    g2_plugin,
+    geom_dataid[dataid_setid, g2],
   )
 
   mesh_data1.nmeshface = nmeshface
@@ -830,7 +850,7 @@ def _sdf_narrowphase(
   mesh_data1.mesh_vert = mesh_vert
   mesh_data1.mesh_faceadr = mesh_faceadr
   mesh_data1.mesh_face = mesh_face
-  mesh_data1.data_id = geom_dataid[g1]
+  mesh_data1.data_id = geom_dataid[dataid_setid, g1]
   mesh_data1.pos = geom1.pos
   mesh_data1.mat = geom1.rot
   mesh_data1.size = geom1.size
@@ -843,7 +863,7 @@ def _sdf_narrowphase(
   mesh_data2.mesh_vert = mesh_vert
   mesh_data2.mesh_faceadr = mesh_faceadr
   mesh_data2.mesh_face = mesh_face
-  mesh_data2.data_id = geom_dataid[g2]
+  mesh_data2.data_id = geom_dataid[dataid_setid, g2]
   mesh_data2.pos = geom2.pos
   mesh_data2.mat = geom2.rot
   mesh_data2.size = geom2.size

mjx/mujoco/mjx/third_party/mujoco_warp/_src/constraint.py

@@ -32,12 +32,12 @@ wp.set_module_options({"enable_backward": False})
 @wp.kernel
 def _zero_constraint_counts(
   # Data out:
-  ne_out: wp.array(dtype=int),
-  nf_out: wp.array(dtype=int),
-  nl_out: wp.array(dtype=int),
-  nefc_out: wp.array(dtype=int),
+  ne_out: wp.array[int],
+  nf_out: wp.array[int],
+  nl_out: wp.array[int],
+  nefc_out: wp.array[int],
   # Out:
-  efc_nnz_out: wp.array(dtype=int),
+  efc_nnz_out: wp.array[int],
 ):
   worldid = wp.tid()
 
@@ -68,14 +68,14 @@ def _efc_row(
   type: int,
   id: int,
   # Out:
-  type_out: wp.array2d(dtype=int),
-  id_out: wp.array2d(dtype=int),
-  pos_out: wp.array2d(dtype=float),
-  margin_out: wp.array2d(dtype=float),
-  D_out: wp.array2d(dtype=float),
-  vel_out: wp.array2d(dtype=float),
-  aref_out: wp.array2d(dtype=float),
-  frictionloss_out: wp.array2d(dtype=float),
+  type_out: wp.array2d[int],
+  id_out: wp.array2d[int],
+  pos_out: wp.array2d[float],
+  margin_out: wp.array2d[float],
+  D_out: wp.array2d[float],
+  vel_out: wp.array2d[float],
+  aref_out: wp.array2d[float],
+  frictionloss_out: wp.array2d[float],
 ):
   # calculate kbi
   timeconst = solref[0]
@@ -126,52 +126,52 @@ def _equality_connect(
   # Model:
   nv: int,
   nsite: int,
-  opt_timestep: wp.array(dtype=float),
+  opt_timestep: wp.array[float],
   opt_disableflags: int,
-  body_parentid: wp.array(dtype=int),
-  body_rootid: wp.array(dtype=int),
-  body_weldid: wp.array(dtype=int),
-  body_dofnum: wp.array(dtype=int),
-  body_dofadr: wp.array(dtype=int),
-  body_invweight0: wp.array2d(dtype=wp.vec2),
-  dof_bodyid: wp.array(dtype=int),
-  dof_parentid: wp.array(dtype=int),
-  site_bodyid: wp.array(dtype=int),
-  eq_obj1id: wp.array(dtype=int),
-  eq_obj2id: wp.array(dtype=int),
-  eq_objtype: wp.array(dtype=int),
-  eq_solref: wp.array2d(dtype=wp.vec2),
-  eq_solimp: wp.array2d(dtype=vec5),
-  eq_data: wp.array2d(dtype=vec11),
+  body_parentid: wp.array[int],
+  body_rootid: wp.array[int],
+  body_weldid: wp.array[int],
+  body_dofnum: wp.array[int],
+  body_dofadr: wp.array[int],
+  body_invweight0: wp.array2d[wp.vec2],
+  dof_bodyid: wp.array[int],
+  dof_parentid: wp.array[int],
+  site_bodyid: wp.array[int],
+  eq_obj1id: wp.array[int],
+  eq_obj2id: wp.array[int],
+  eq_objtype: wp.array[int],
+  eq_solref: wp.array2d[wp.vec2],
+  eq_solimp: wp.array2d[vec5],
+  eq_data: wp.array2d[vec11],
   is_sparse: bool,
-  eq_connect_adr: wp.array(dtype=int),
+  eq_connect_adr: wp.array[int],
   # Data in:
-  qvel_in: wp.array2d(dtype=float),
-  eq_active_in: wp.array2d(dtype=bool),
-  xpos_in: wp.array2d(dtype=wp.vec3),
-  xmat_in: wp.array2d(dtype=wp.mat33),
-  site_xpos_in: wp.array2d(dtype=wp.vec3),
-  subtree_com_in: wp.array2d(dtype=wp.vec3),
-  cdof_in: wp.array2d(dtype=wp.spatial_vector),
+  qvel_in: wp.array2d[float],
+  eq_active_in: wp.array2d[bool],
+  xpos_in: wp.array2d[wp.vec3],
+  xmat_in: wp.array2d[wp.mat33],
+  site_xpos_in: wp.array2d[wp.vec3],
+  subtree_com_in: wp.array2d[wp.vec3],
+  cdof_in: wp.array2d[wp.spatial_vector],
   njmax_in: int,
   njmax_nnz_in: int,
   # Data out:
-  ne_out: wp.array(dtype=int),
-  nefc_out: wp.array(dtype=int),
-  efc_type_out: wp.array2d(dtype=int),
-  efc_id_out: wp.array2d(dtype=int),
-  efc_J_rownnz_out: wp.array2d(dtype=int),
-  efc_J_rowadr_out: wp.array2d(dtype=int),
-  efc_J_colind_out: wp.array3d(dtype=int),
-  efc_J_out: wp.array3d(dtype=float),
-  efc_pos_out: wp.array2d(dtype=float),
-  efc_margin_out: wp.array2d(dtype=float),
-  efc_D_out: wp.array2d(dtype=float),
-  efc_vel_out: wp.array2d(dtype=float),
-  efc_aref_out: wp.array2d(dtype=float),
-  efc_frictionloss_out: wp.array2d(dtype=float),
+  ne_out: wp.array[int],
+  nefc_out: wp.array[int],
+  efc_type_out: wp.array2d[int],
+  efc_id_out: wp.array2d[int],
+  efc_J_rownnz_out: wp.array2d[int],
+  efc_J_rowadr_out: wp.array2d[int],
+  efc_J_colind_out: wp.array3d[int],
+  efc_J_out: wp.array3d[float],
+  efc_pos_out: wp.array2d[float],
+  efc_margin_out: wp.array2d[float],
+  efc_D_out: wp.array2d[float],
+  efc_vel_out: wp.array2d[float],
+  efc_aref_out: wp.array2d[float],
+  efc_frictionloss_out: wp.array2d[float],
   # Out:
-  efc_nnz_out: wp.array(dtype=int),
+  efc_nnz_out: wp.array[int],
 ):
   """Calculates constraint rows for connect equality constraints."""
   worldid, eqconnectid = wp.tid()
@@ -368,42 +368,42 @@ def _equality_connect(
 def _equality_joint(
   # Model:
   nv: int,
-  opt_timestep: wp.array(dtype=float),
+  opt_timestep: wp.array[float],
   opt_disableflags: int,
-  qpos0: wp.array2d(dtype=float),
-  jnt_qposadr: wp.array(dtype=int),
-  jnt_dofadr: wp.array(dtype=int),
-  dof_invweight0: wp.array2d(dtype=float),
-  eq_obj1id: wp.array(dtype=int),
-  eq_obj2id: wp.array(dtype=int),
-  eq_solref: wp.array2d(dtype=wp.vec2),
-  eq_solimp: wp.array2d(dtype=vec5),
-  eq_data: wp.array2d(dtype=vec11),
+  qpos0: wp.array2d[float],
+  jnt_qposadr: wp.array[int],
+  jnt_dofadr: wp.array[int],
+  dof_invweight0: wp.array2d[float],
+  eq_obj1id: wp.array[int],
+  eq_obj2id: wp.array[int],
+  eq_solref: wp.array2d[wp.vec2],
+  eq_solimp: wp.array2d[vec5],
+  eq_data: wp.array2d[vec11],
   is_sparse: bool,
-  eq_jnt_adr: wp.array(dtype=int),
+  eq_jnt_adr: wp.array[int],
   # Data in:
-  qpos_in: wp.array2d(dtype=float),
-  qvel_in: wp.array2d(dtype=float),
-  eq_active_in: wp.array2d(dtype=bool),
+  qpos_in: wp.array2d[float],
+  qvel_in: wp.array2d[float],
+  eq_active_in: wp.array2d[bool],
   njmax_in: int,
   njmax_nnz_in: int,
   # Data out:
-  ne_out: wp.array(dtype=int),
-  nefc_out: wp.array(dtype=int),
-  efc_type_out: wp.array2d(dtype=int),
-  efc_id_out: wp.array2d(dtype=int),
-  efc_J_rownnz_out: wp.array2d(dtype=int),
-  efc_J_rowadr_out: wp.array2d(dtype=int),
-  efc_J_colind_out: wp.array3d(dtype=int),
-  efc_J_out: wp.array3d(dtype=float),
-  efc_pos_out: wp.array2d(dtype=float),
-  efc_margin_out: wp.array2d(dtype=float),
-  efc_D_out: wp.array2d(dtype=float),
-  efc_vel_out: wp.array2d(dtype=float),
-  efc_aref_out: wp.array2d(dtype=float),
-  efc_frictionloss_out: wp.array2d(dtype=float),
+  ne_out: wp.array[int],
+  nefc_out: wp.array[int],
+  efc_type_out: wp.array2d[int],
+  efc_id_out: wp.array2d[int],
+  efc_J_rownnz_out: wp.array2d[int],
+  efc_J_rowadr_out: wp.array2d[int],
+  efc_J_colind_out: wp.array3d[int],
+  efc_J_out: wp.array3d[float],
+  efc_pos_out: wp.array2d[float],
+  efc_margin_out: wp.array2d[float],
+  efc_D_out: wp.array2d[float],
+  efc_vel_out: wp.array2d[float],
+  efc_aref_out: wp.array2d[float],
+  efc_frictionloss_out: wp.array2d[float],
   # Out:
-  efc_nnz_out: wp.array(dtype=int),
+  efc_nnz_out: wp.array[int],
 ):
   worldid, eqjntid = wp.tid()
   eqid = eq_jnt_adr[eqjntid]
@@ -499,44 +499,44 @@ def _equality_joint(
 def _equality_tendon(
   # Model:
   nv: int,
-  opt_timestep: wp.array(dtype=float),
+  opt_timestep: wp.array[float],
   opt_disableflags: int,
-  eq_obj1id: wp.array(dtype=int),
-  eq_obj2id: wp.array(dtype=int),
-  eq_solref: wp.array2d(dtype=wp.vec2),
-  eq_solimp: wp.array2d(dtype=vec5),
-  eq_data: wp.array2d(dtype=vec11),
-  ten_J_rownnz: wp.array(dtype=int),
-  ten_J_rowadr: wp.array(dtype=int),
-  ten_J_colind: wp.array(dtype=int),
-  tendon_length0: wp.array2d(dtype=float),
-  tendon_invweight0: wp.array2d(dtype=float),
+  eq_obj1id: wp.array[int],
+  eq_obj2id: wp.array[int],
+  eq_solref: wp.array2d[wp.vec2],
+  eq_solimp: wp.array2d[vec5],
+  eq_data: wp.array2d[vec11],
+  ten_J_rownnz: wp.array[int],
+  ten_J_rowadr: wp.array[int],
+  ten_J_colind: wp.array[int],
+  tendon_length0: wp.array2d[float],
+  tendon_invweight0: wp.array2d[float],
   is_sparse: bool,
-  eq_ten_adr: wp.array(dtype=int),
+  eq_ten_adr: wp.array[int],
   # Data in:
-  qvel_in: wp.array2d(dtype=float),
-  eq_active_in: wp.array2d(dtype=bool),
-  ten_J_in: wp.array2d(dtype=float),
-  ten_length_in: wp.array2d(dtype=float),
+  qvel_in: wp.array2d[float],
+  eq_active_in: wp.array2d[bool],
+  ten_J_in: wp.array2d[float],
+  ten_length_in: wp.array2d[float],
   njmax_in: int,
   njmax_nnz_in: int,
   # Data out:
-  ne_out: wp.array(dtype=int),
-  nefc_out: wp.array(dtype=int),
-  efc_type_out: wp.array2d(dtype=int),
-  efc_id_out: wp.array2d(dtype=int),
-  efc_J_rownnz_out: wp.array2d(dtype=int),
-  efc_J_rowadr_out: wp.array2d(dtype=int),
-  efc_J_colind_out: wp.array3d(dtype=int),
-  efc_J_out: wp.array3d(dtype=float),
-  efc_pos_out: wp.array2d(dtype=float),
-  efc_margin_out: wp.array2d(dtype=float),
-  efc_D_out: wp.array2d(dtype=float),
-  efc_vel_out: wp.array2d(dtype=float),
-  efc_aref_out: wp.array2d(dtype=float),
-  efc_frictionloss_out: wp.array2d(dtype=float),
+  ne_out: wp.array[int],
+  nefc_out: wp.array[int],
+  efc_type_out: wp.array2d[int],
+  efc_id_out: wp.array2d[int],
+  efc_J_rownnz_out: wp.array2d[int],
+  efc_J_rowadr_out: wp.array2d[int],
+  efc_J_colind_out: wp.array3d[int],
+  efc_J_out: wp.array3d[float],
+  efc_pos_out: wp.array2d[float],
+  efc_margin_out: wp.array2d[float],
+  efc_D_out: wp.array2d[float],
+  efc_vel_out: wp.array2d[float],
+  efc_aref_out: wp.array2d[float],
+  efc_frictionloss_out: wp.array2d[float],
   # Out:
-  efc_nnz_out: wp.array(dtype=int),
+  efc_nnz_out: wp.array[int],
 ):
   worldid, eqtenid = wp.tid()
   eqid = eq_ten_adr[eqtenid]
@@ -679,42 +679,42 @@ def _equality_flex(is_sparse: bool):
   def kernel(
     # Model:
     nv: int,
-    opt_timestep: wp.array(dtype=float),
+    opt_timestep: wp.array[float],
     opt_disableflags: int,
-    flex_edgeadr: wp.array(dtype=int),
-    flex_edgenum: wp.array(dtype=int),
-    flexedge_length0: wp.array(dtype=float),
-    flexedge_invweight0: wp.array(dtype=float),
-    flexedge_J_rownnz: wp.array(dtype=int),
-    flexedge_J_rowadr: wp.array(dtype=int),
-    flexedge_J_colind: wp.array(dtype=int),
-    eq_obj1id: wp.array(dtype=int),
-    eq_solref: wp.array2d(dtype=wp.vec2),
-    eq_solimp: wp.array2d(dtype=vec5),
-    eq_flex_adr: wp.array(dtype=int),
+    flex_edgeadr: wp.array[int],
+    flex_edgenum: wp.array[int],
+    flexedge_length0: wp.array[float],
+    flexedge_invweight0: wp.array[float],
+    flexedge_J_rownnz: wp.array[int],
+    flexedge_J_rowadr: wp.array[int],
+    flexedge_J_colind: wp.array[int],
+    eq_obj1id: wp.array[int],
+    eq_solref: wp.array2d[wp.vec2],
+    eq_solimp: wp.array2d[vec5],
+    eq_flex_adr: wp.array[int],
     # Data in:
-    qvel_in: wp.array2d(dtype=float),
-    flexedge_J_in: wp.array2d(dtype=float),
-    flexedge_length_in: wp.array2d(dtype=float),
+    qvel_in: wp.array2d[float],
+    flexedge_J_in: wp.array2d[float],
+    flexedge_length_in: wp.array2d[float],
     njmax_in: int,
     njmax_nnz_in: int,
     # Data out:
-    ne_out: wp.array(dtype=int),
-    nefc_out: wp.array(dtype=int),
-    efc_type_out: wp.array2d(dtype=int),
-    efc_id_out: wp.array2d(dtype=int),
-    efc_J_rownnz_out: wp.array2d(dtype=int),
-    efc_J_rowadr_out: wp.array2d(dtype=int),
-    efc_J_colind_out: wp.array3d(dtype=int),
-    efc_J_out: wp.array3d(dtype=float),
-    efc_pos_out: wp.array2d(dtype=float),
-    efc_margin_out: wp.array2d(dtype=float),
-    efc_D_out: wp.array2d(dtype=float),
-    efc_vel_out: wp.array2d(dtype=float),
-    efc_aref_out: wp.array2d(dtype=float),
-    efc_frictionloss_out: wp.array2d(dtype=float),
+    ne_out: wp.array[int],
+    nefc_out: wp.array[int],
+    efc_type_out: wp.array2d[int],
+    efc_id_out: wp.array2d[int],
+    efc_J_rownnz_out: wp.array2d[int],
+    efc_J_rowadr_out: wp.array2d[int],
+    efc_J_colind_out: wp.array3d[int],
+    efc_J_out: wp.array3d[float],
+    efc_pos_out: wp.array2d[float],
+    efc_margin_out: wp.array2d[float],
+    efc_D_out: wp.array2d[float],
+    efc_vel_out: wp.array2d[float],
+    efc_aref_out: wp.array2d[float],
+    efc_frictionloss_out: wp.array2d[float],
     # Out:
-    efc_nnz_out: wp.array(dtype=int),
+    efc_nnz_out: wp.array[int],
   ):
     worldid, eqflexid, edgeid = wp.tid()
     eqid = eq_flex_adr[eqflexid]
@@ -794,54 +794,54 @@ def _equality_weld(
   # Model:
   nv: int,
   nsite: int,
-  opt_timestep: wp.array(dtype=float),
+  opt_timestep: wp.array[float],
   opt_disableflags: int,
-  body_parentid: wp.array(dtype=int),
-  body_rootid: wp.array(dtype=int),
-  body_weldid: wp.array(dtype=int),
-  body_dofnum: wp.array(dtype=int),
-  body_dofadr: wp.array(dtype=int),
-  body_invweight0: wp.array2d(dtype=wp.vec2),
-  dof_bodyid: wp.array(dtype=int),
-  dof_parentid: wp.array(dtype=int),
-  site_bodyid: wp.array(dtype=int),
-  site_quat: wp.array2d(dtype=wp.quat),
-  eq_obj1id: wp.array(dtype=int),
-  eq_obj2id: wp.array(dtype=int),
-  eq_objtype: wp.array(dtype=int),
-  eq_solref: wp.array2d(dtype=wp.vec2),
-  eq_solimp: wp.array2d(dtype=vec5),
-  eq_data: wp.array2d(dtype=vec11),
+  body_parentid: wp.array[int],
+  body_rootid: wp.array[int],
+  body_weldid: wp.array[int],
+  body_dofnum: wp.array[int],
+  body_dofadr: wp.array[int],
+  body_invweight0: wp.array2d[wp.vec2],
+  dof_bodyid: wp.array[int],
+  dof_parentid: wp.array[int],
+  site_bodyid: wp.array[int],
+  site_quat: wp.array2d[wp.quat],
+  eq_obj1id: wp.array[int],
+  eq_obj2id: wp.array[int],
+  eq_objtype: wp.array[int],
+  eq_solref: wp.array2d[wp.vec2],
+  eq_solimp: wp.array2d[vec5],
+  eq_data: wp.array2d[vec11],
   is_sparse: bool,
-  eq_wld_adr: wp.array(dtype=int),
+  eq_wld_adr: wp.array[int],
   # Data in:
-  qvel_in: wp.array2d(dtype=float),
-  eq_active_in: wp.array2d(dtype=bool),
-  xpos_in: wp.array2d(dtype=wp.vec3),
-  xquat_in: wp.array2d(dtype=wp.quat),
-  xmat_in: wp.array2d(dtype=wp.mat33),
-  site_xpos_in: wp.array2d(dtype=wp.vec3),
-  subtree_com_in: wp.array2d(dtype=wp.vec3),
-  cdof_in: wp.array2d(dtype=wp.spatial_vector),
+  qvel_in: wp.array2d[float],
+  eq_active_in: wp.array2d[bool],
+  xpos_in: wp.array2d[wp.vec3],
+  xquat_in: wp.array2d[wp.quat],
+  xmat_in: wp.array2d[wp.mat33],
+  site_xpos_in: wp.array2d[wp.vec3],
+  subtree_com_in: wp.array2d[wp.vec3],
+  cdof_in: wp.array2d[wp.spatial_vector],
   njmax_in: int,
   njmax_nnz_in: int,
   # Data out:
-  ne_out: wp.array(dtype=int),
-  nefc_out: wp.array(dtype=int),
-  efc_type_out: wp.array2d(dtype=int),
-  efc_id_out: wp.array2d(dtype=int),
-  efc_J_rownnz_out: wp.array2d(dtype=int),
-  efc_J_rowadr_out: wp.array2d(dtype=int),
-  efc_J_colind_out: wp.array3d(dtype=int),
-  efc_J_out: wp.array3d(dtype=float),
-  efc_pos_out: wp.array2d(dtype=float),
-  efc_margin_out: wp.array2d(dtype=float),
-  efc_D_out: wp.array2d(dtype=float),
-  efc_vel_out: wp.array2d(dtype=float),
-  efc_aref_out: wp.array2d(dtype=float),
-  efc_frictionloss_out: wp.array2d(dtype=float),
+  ne_out: wp.array[int],
+  nefc_out: wp.array[int],
+  efc_type_out: wp.array2d[int],
+  efc_id_out: wp.array2d[int],
+  efc_J_rownnz_out: wp.array2d[int],
+  efc_J_rowadr_out: wp.array2d[int],
+  efc_J_colind_out: wp.array3d[int],
+  efc_J_out: wp.array3d[float],
+  efc_pos_out: wp.array2d[float],
+  efc_margin_out: wp.array2d[float],
+  efc_D_out: wp.array2d[float],
+  efc_vel_out: wp.array2d[float],
+  efc_aref_out: wp.array2d[float],
+  efc_frictionloss_out: wp.array2d[float],
   # Out:
-  efc_nnz_out: wp.array(dtype=int),
+  efc_nnz_out: wp.array[int],
 ):
   worldid, eqweldid = wp.tid()
   eqid = eq_wld_adr[eqweldid]
@@ -1114,34 +1114,34 @@ def _equality_weld(
 def _friction_dof(
   # Model:
   nv: int,
-  opt_timestep: wp.array(dtype=float),
+  opt_timestep: wp.array[float],
   opt_disableflags: int,
-  dof_solref: wp.array2d(dtype=wp.vec2),
-  dof_solimp: wp.array2d(dtype=vec5),
-  dof_frictionloss: wp.array2d(dtype=float),
-  dof_invweight0: wp.array2d(dtype=float),
+  dof_solref: wp.array2d[wp.vec2],
+  dof_solimp: wp.array2d[vec5],
+  dof_frictionloss: wp.array2d[float],
+  dof_invweight0: wp.array2d[float],
   is_sparse: bool,
   # Data in:
-  qvel_in: wp.array2d(dtype=float),
+  qvel_in: wp.array2d[float],
   njmax_in: int,
   njmax_nnz_in: int,
   # Data out:
-  nf_out: wp.array(dtype=int),
-  nefc_out: wp.array(dtype=int),
-  efc_type_out: wp.array2d(dtype=int),
-  efc_id_out: wp.array2d(dtype=int),
-  efc_J_rownnz_out: wp.array2d(dtype=int),
-  efc_J_rowadr_out: wp.array2d(dtype=int),
-  efc_J_colind_out: wp.array3d(dtype=int),
-  efc_J_out: wp.array3d(dtype=float),
-  efc_pos_out: wp.array2d(dtype=float),
-  efc_margin_out: wp.array2d(dtype=float),
-  efc_D_out: wp.array2d(dtype=float),
-  efc_vel_out: wp.array2d(dtype=float),
-  efc_aref_out: wp.array2d(dtype=float),
-  efc_frictionloss_out: wp.array2d(dtype=float),
+  nf_out: wp.array[int],
+  nefc_out: wp.array[int],
+  efc_type_out: wp.array2d[int],
+  efc_id_out: wp.array2d[int],
+  efc_J_rownnz_out: wp.array2d[int],
+  efc_J_rowadr_out: wp.array2d[int],
+  efc_J_colind_out: wp.array3d[int],
+  efc_J_out: wp.array3d[float],
+  efc_pos_out: wp.array2d[float],
+  efc_margin_out: wp.array2d[float],
+  efc_D_out: wp.array2d[float],
+  efc_vel_out: wp.array2d[float],
+  efc_aref_out: wp.array2d[float],
+  efc_frictionloss_out: wp.array2d[float],
   # Out:
-  efc_nnz_out: wp.array(dtype=int),
+  efc_nnz_out: wp.array[int],
 ):
   worldid, dofid = wp.tid()
 
@@ -1204,38 +1204,38 @@ def _friction_dof(
 def _friction_tendon(
   # Model:
   nv: int,
-  opt_timestep: wp.array(dtype=float),
+  opt_timestep: wp.array[float],
   opt_disableflags: int,
-  ten_J_rownnz: wp.array(dtype=int),
-  ten_J_rowadr: wp.array(dtype=int),
-  ten_J_colind: wp.array(dtype=int),
-  tendon_solref_fri: wp.array2d(dtype=wp.vec2),
-  tendon_solimp_fri: wp.array2d(dtype=vec5),
-  tendon_frictionloss: wp.array2d(dtype=float),
-  tendon_invweight0: wp.array2d(dtype=float),
+  ten_J_rownnz: wp.array[int],
+  ten_J_rowadr: wp.array[int],
+  ten_J_colind: wp.array[int],
+  tendon_solref_fri: wp.array2d[wp.vec2],
+  tendon_solimp_fri: wp.array2d[vec5],
+  tendon_frictionloss: wp.array2d[float],
+  tendon_invweight0: wp.array2d[float],
   is_sparse: bool,
   # Data in:
-  qvel_in: wp.array2d(dtype=float),
-  ten_J_in: wp.array2d(dtype=float),
+  qvel_in: wp.array2d[float],
+  ten_J_in: wp.array2d[float],
   njmax_in: int,
   njmax_nnz_in: int,
   # Data out:
-  nf_out: wp.array(dtype=int),
-  nefc_out: wp.array(dtype=int),
-  efc_type_out: wp.array2d(dtype=int),
-  efc_id_out: wp.array2d(dtype=int),
-  efc_J_rownnz_out: wp.array2d(dtype=int),
-  efc_J_rowadr_out: wp.array2d(dtype=int),
-  efc_J_colind_out: wp.array3d(dtype=int),
-  efc_J_out: wp.array3d(dtype=float),
-  efc_pos_out: wp.array2d(dtype=float),
-  efc_margin_out: wp.array2d(dtype=float),
-  efc_D_out: wp.array2d(dtype=float),
-  efc_vel_out: wp.array2d(dtype=float),
-  efc_aref_out: wp.array2d(dtype=float),
-  efc_frictionloss_out: wp.array2d(dtype=float),
+  nf_out: wp.array[int],
+  nefc_out: wp.array[int],
+  efc_type_out: wp.array2d[int],
+  efc_id_out: wp.array2d[int],
+  efc_J_rownnz_out: wp.array2d[int],
+  efc_J_rowadr_out: wp.array2d[int],
+  efc_J_colind_out: wp.array3d[int],
+  efc_J_out: wp.array3d[float],
+  efc_pos_out: wp.array2d[float],
+  efc_margin_out: wp.array2d[float],
+  efc_D_out: wp.array2d[float],
+  efc_vel_out: wp.array2d[float],
+  efc_aref_out: wp.array2d[float],
+  efc_frictionloss_out: wp.array2d[float],
   # Out:
-  efc_nnz_out: wp.array(dtype=int),
+  efc_nnz_out: wp.array[int],
 ):
   worldid, tenid = wp.tid()
 
@@ -1317,39 +1317,39 @@ def _friction_tendon(
 def _limit_slide_hinge(
   # Model:
   nv: int,
-  opt_timestep: wp.array(dtype=float),
+  opt_timestep: wp.array[float],
   opt_disableflags: int,
-  jnt_qposadr: wp.array(dtype=int),
-  jnt_dofadr: wp.array(dtype=int),
-  jnt_solref: wp.array2d(dtype=wp.vec2),
-  jnt_solimp: wp.array2d(dtype=vec5),
-  jnt_range: wp.array2d(dtype=wp.vec2),
-  jnt_margin: wp.array2d(dtype=float),
-  dof_invweight0: wp.array2d(dtype=float),
+  jnt_qposadr: wp.array[int],
+  jnt_dofadr: wp.array[int],
+  jnt_solref: wp.array2d[wp.vec2],
+  jnt_solimp: wp.array2d[vec5],
+  jnt_range: wp.array2d[wp.vec2],
+  jnt_margin: wp.array2d[float],
+  dof_invweight0: wp.array2d[float],
   is_sparse: bool,
-  jnt_limited_slide_hinge_adr: wp.array(dtype=int),
+  jnt_limited_slide_hinge_adr: wp.array[int],
   # Data in:
-  qpos_in: wp.array2d(dtype=float),
-  qvel_in: wp.array2d(dtype=float),
+  qpos_in: wp.array2d[float],
+  qvel_in: wp.array2d[float],
   njmax_in: int,
   njmax_nnz_in: int,
   # Data out:
-  nl_out: wp.array(dtype=int),
-  nefc_out: wp.array(dtype=int),
-  efc_type_out: wp.array2d(dtype=int),
-  efc_id_out: wp.array2d(dtype=int),
-  efc_J_rownnz_out: wp.array2d(dtype=int),
-  efc_J_rowadr_out: wp.array2d(dtype=int),
-  efc_J_colind_out: wp.array3d(dtype=int),
-  efc_J_out: wp.array3d(dtype=float),
-  efc_pos_out: wp.array2d(dtype=float),
-  efc_margin_out: wp.array2d(dtype=float),
-  efc_D_out: wp.array2d(dtype=float),
-  efc_vel_out: wp.array2d(dtype=float),
-  efc_aref_out: wp.array2d(dtype=float),
-  efc_frictionloss_out: wp.array2d(dtype=float),
+  nl_out: wp.array[int],
+  nefc_out: wp.array[int],
+  efc_type_out: wp.array2d[int],
+  efc_id_out: wp.array2d[int],
+  efc_J_rownnz_out: wp.array2d[int],
+  efc_J_rowadr_out: wp.array2d[int],
+  efc_J_colind_out: wp.array3d[int],
+  efc_J_out: wp.array3d[float],
+  efc_pos_out: wp.array2d[float],
+  efc_margin_out: wp.array2d[float],
+  efc_D_out: wp.array2d[float],
+  efc_vel_out: wp.array2d[float],
+  efc_aref_out: wp.array2d[float],
+  efc_frictionloss_out: wp.array2d[float],
   # Out:
-  efc_nnz_out: wp.array(dtype=int),
+  efc_nnz_out: wp.array[int],
 ):
   worldid, jntlimitedid = wp.tid()
   jntid = jnt_limited_slide_hinge_adr[jntlimitedid]
@@ -1422,39 +1422,39 @@ def _limit_slide_hinge(
 def _limit_ball(
   # Model:
   nv: int,
-  opt_timestep: wp.array(dtype=float),
+  opt_timestep: wp.array[float],
   opt_disableflags: int,
-  jnt_qposadr: wp.array(dtype=int),
-  jnt_dofadr: wp.array(dtype=int),
-  jnt_solref: wp.array2d(dtype=wp.vec2),
-  jnt_solimp: wp.array2d(dtype=vec5),
-  jnt_range: wp.array2d(dtype=wp.vec2),
-  jnt_margin: wp.array2d(dtype=float),
-  dof_invweight0: wp.array2d(dtype=float),
+  jnt_qposadr: wp.array[int],
+  jnt_dofadr: wp.array[int],
+  jnt_solref: wp.array2d[wp.vec2],
+  jnt_solimp: wp.array2d[vec5],
+  jnt_range: wp.array2d[wp.vec2],
+  jnt_margin: wp.array2d[float],
+  dof_invweight0: wp.array2d[float],
   is_sparse: bool,
-  jnt_limited_ball_adr: wp.array(dtype=int),
+  jnt_limited_ball_adr: wp.array[int],
   # Data in:
-  qpos_in: wp.array2d(dtype=float),
-  qvel_in: wp.array2d(dtype=float),
+  qpos_in: wp.array2d[float],
+  qvel_in: wp.array2d[float],
   njmax_in: int,
   njmax_nnz_in: int,
   # Data out:
-  nl_out: wp.array(dtype=int),
-  nefc_out: wp.array(dtype=int),
-  efc_type_out: wp.array2d(dtype=int),
-  efc_id_out: wp.array2d(dtype=int),
-  efc_J_rownnz_out: wp.array2d(dtype=int),
-  efc_J_rowadr_out: wp.array2d(dtype=int),
-  efc_J_colind_out: wp.array3d(dtype=int),
-  efc_J_out: wp.array3d(dtype=float),
-  efc_pos_out: wp.array2d(dtype=float),
-  efc_margin_out: wp.array2d(dtype=float),
-  efc_D_out: wp.array2d(dtype=float),
-  efc_vel_out: wp.array2d(dtype=float),
-  efc_aref_out: wp.array2d(dtype=float),
-  efc_frictionloss_out: wp.array2d(dtype=float),
+  nl_out: wp.array[int],
+  nefc_out: wp.array[int],
+  efc_type_out: wp.array2d[int],
+  efc_id_out: wp.array2d[int],
+  efc_J_rownnz_out: wp.array2d[int],
+  efc_J_rowadr_out: wp.array2d[int],
+  efc_J_colind_out: wp.array3d[int],
+  efc_J_out: wp.array3d[float],
+  efc_pos_out: wp.array2d[float],
+  efc_margin_out: wp.array2d[float],
+  efc_D_out: wp.array2d[float],
+  efc_vel_out: wp.array2d[float],
+  efc_aref_out: wp.array2d[float],
+  efc_frictionloss_out: wp.array2d[float],
   # Out:
-  efc_nnz_out: wp.array(dtype=int),
+  efc_nnz_out: wp.array[int],
 ):
   worldid, jntlimitedid = wp.tid()
   jntid = jnt_limited_ball_adr[jntlimitedid]
@@ -1547,41 +1547,41 @@ def _limit_ball(
 def _limit_tendon(
   # Model:
   nv: int,
-  opt_timestep: wp.array(dtype=float),
+  opt_timestep: wp.array[float],
   opt_disableflags: int,
-  ten_J_rownnz: wp.array(dtype=int),
-  ten_J_rowadr: wp.array(dtype=int),
-  ten_J_colind: wp.array(dtype=int),
-  tendon_solref_lim: wp.array2d(dtype=wp.vec2),
-  tendon_solimp_lim: wp.array2d(dtype=vec5),
-  tendon_range: wp.array2d(dtype=wp.vec2),
-  tendon_margin: wp.array2d(dtype=float),
-  tendon_invweight0: wp.array2d(dtype=float),
+  ten_J_rownnz: wp.array[int],
+  ten_J_rowadr: wp.array[int],
+  ten_J_colind: wp.array[int],
+  tendon_solref_lim: wp.array2d[wp.vec2],
+  tendon_solimp_lim: wp.array2d[vec5],
+  tendon_range: wp.array2d[wp.vec2],
+  tendon_margin: wp.array2d[float],
+  tendon_invweight0: wp.array2d[float],
   is_sparse: bool,
-  tendon_limited_adr: wp.array(dtype=int),
+  tendon_limited_adr: wp.array[int],
   # Data in:
-  qvel_in: wp.array2d(dtype=float),
-  ten_J_in: wp.array2d(dtype=float),
-  ten_length_in: wp.array2d(dtype=float),
+  qvel_in: wp.array2d[float],
+  ten_J_in: wp.array2d[float],
+  ten_length_in: wp.array2d[float],
   njmax_in: int,
   njmax_nnz_in: int,
   # Data out:
-  nl_out: wp.array(dtype=int),
-  nefc_out: wp.array(dtype=int),
-  efc_type_out: wp.array2d(dtype=int),
-  efc_id_out: wp.array2d(dtype=int),
-  efc_J_rownnz_out: wp.array2d(dtype=int),
-  efc_J_rowadr_out: wp.array2d(dtype=int),
-  efc_J_colind_out: wp.array3d(dtype=int),
-  efc_J_out: wp.array3d(dtype=float),
-  efc_pos_out: wp.array2d(dtype=float),
-  efc_margin_out: wp.array2d(dtype=float),
-  efc_D_out: wp.array2d(dtype=float),
-  efc_vel_out: wp.array2d(dtype=float),
-  efc_aref_out: wp.array2d(dtype=float),
-  efc_frictionloss_out: wp.array2d(dtype=float),
+  nl_out: wp.array[int],
+  nefc_out: wp.array[int],
+  efc_type_out: wp.array2d[int],
+  efc_id_out: wp.array2d[int],
+  efc_J_rownnz_out: wp.array2d[int],
+  efc_J_rowadr_out: wp.array2d[int],
+  efc_J_colind_out: wp.array3d[int],
+  efc_J_out: wp.array3d[float],
+  efc_pos_out: wp.array2d[float],
+  efc_margin_out: wp.array2d[float],
+  efc_D_out: wp.array2d[float],
+  efc_vel_out: wp.array2d[float],
+  efc_aref_out: wp.array2d[float],
+  efc_frictionloss_out: wp.array2d[float],
   # Out:
-  efc_nnz_out: wp.array(dtype=int),
+  efc_nnz_out: wp.array[int],
 ):
   worldid, tenlimitedid = wp.tid()
   tenid = tendon_limited_adr[tenlimitedid]
@@ -1669,59 +1669,59 @@ def _limit_tendon(
 def _contact_pyramidal(
   # Model:
   nv: int,
-  opt_timestep: wp.array(dtype=float),
+  opt_timestep: wp.array[float],
   opt_disableflags: int,
-  opt_impratio_invsqrt: wp.array(dtype=float),
-  body_parentid: wp.array(dtype=int),
-  body_rootid: wp.array(dtype=int),
-  body_weldid: wp.array(dtype=int),
-  body_dofnum: wp.array(dtype=int),
-  body_dofadr: wp.array(dtype=int),
-  body_invweight0: wp.array2d(dtype=wp.vec2),
-  dof_bodyid: wp.array(dtype=int),
-  dof_parentid: wp.array(dtype=int),
-  geom_bodyid: wp.array(dtype=int),
-  flex_vertadr: wp.array(dtype=int),
-  flex_vertbodyid: wp.array(dtype=int),
+  opt_impratio_invsqrt: wp.array[float],
+  body_parentid: wp.array[int],
+  body_rootid: wp.array[int],
+  body_weldid: wp.array[int],
+  body_dofnum: wp.array[int],
+  body_dofadr: wp.array[int],
+  body_invweight0: wp.array2d[wp.vec2],
+  dof_bodyid: wp.array[int],
+  dof_parentid: wp.array[int],
+  geom_bodyid: wp.array[int],
+  flex_vertadr: wp.array[int],
+  flex_vertbodyid: wp.array[int],
   is_sparse: bool,
   # Data in:
-  qvel_in: wp.array2d(dtype=float),
-  subtree_com_in: wp.array2d(dtype=wp.vec3),
-  cdof_in: wp.array2d(dtype=wp.spatial_vector),
+  qvel_in: wp.array2d[float],
+  subtree_com_in: wp.array2d[wp.vec3],
+  cdof_in: wp.array2d[wp.spatial_vector],
   njmax_in: int,
   njmax_nnz_in: int,
-  nacon_in: wp.array(dtype=int),
+  nacon_in: wp.array[int],
   # In:
-  dist_in: wp.array(dtype=float),
-  condim_in: wp.array(dtype=int),
-  includemargin_in: wp.array(dtype=float),
-  worldid_in: wp.array(dtype=int),
-  geom_in: wp.array(dtype=wp.vec2i),
-  flex_in: wp.array(dtype=wp.vec2i),
-  vert_in: wp.array(dtype=wp.vec2i),
-  pos_in: wp.array(dtype=wp.vec3),
-  frame_in: wp.array(dtype=wp.mat33),
-  friction_in: wp.array(dtype=vec5),
-  solref_in: wp.array(dtype=wp.vec2),
-  solimp_in: wp.array(dtype=vec5),
-  type_in: wp.array(dtype=int),
+  dist_in: wp.array[float],
+  condim_in: wp.array[int],
+  includemargin_in: wp.array[float],
+  worldid_in: wp.array[int],
+  geom_in: wp.array[wp.vec2i],
+  flex_in: wp.array[wp.vec2i],
+  vert_in: wp.array[wp.vec2i],
+  pos_in: wp.array[wp.vec3],
+  frame_in: wp.array[wp.mat33],
+  friction_in: wp.array[vec5],
+  solref_in: wp.array[wp.vec2],
+  solimp_in: wp.array[vec5],
+  type_in: wp.array[int],
   # Data out:
-  nefc_out: wp.array(dtype=int),
-  contact_efc_address_out: wp.array2d(dtype=int),
-  efc_type_out: wp.array2d(dtype=int),
-  efc_id_out: wp.array2d(dtype=int),
-  efc_J_rownnz_out: wp.array2d(dtype=int),
-  efc_J_rowadr_out: wp.array2d(dtype=int),
-  efc_J_colind_out: wp.array3d(dtype=int),
-  efc_J_out: wp.array3d(dtype=float),
-  efc_pos_out: wp.array2d(dtype=float),
-  efc_margin_out: wp.array2d(dtype=float),
-  efc_D_out: wp.array2d(dtype=float),
-  efc_vel_out: wp.array2d(dtype=float),
-  efc_aref_out: wp.array2d(dtype=float),
-  efc_frictionloss_out: wp.array2d(dtype=float),
+  nefc_out: wp.array[int],
+  contact_efc_address_out: wp.array2d[int],
+  efc_type_out: wp.array2d[int],
+  efc_id_out: wp.array2d[int],
+  efc_J_rownnz_out: wp.array2d[int],
+  efc_J_rowadr_out: wp.array2d[int],
+  efc_J_colind_out: wp.array3d[int],
+  efc_J_out: wp.array3d[float],
+  efc_pos_out: wp.array2d[float],
+  efc_margin_out: wp.array2d[float],
+  efc_D_out: wp.array2d[float],
+  efc_vel_out: wp.array2d[float],
+  efc_aref_out: wp.array2d[float],
+  efc_frictionloss_out: wp.array2d[float],
   # Out:
-  efc_nnz_out: wp.array(dtype=int),
+  efc_nnz_out: wp.array[int],
 ):
   conid, dimid = wp.tid()
 
@@ -1940,60 +1940,60 @@ def _contact_pyramidal(
 def _contact_elliptic(
   # Model:
   nv: int,
-  opt_timestep: wp.array(dtype=float),
+  opt_timestep: wp.array[float],
   opt_disableflags: int,
-  opt_impratio_invsqrt: wp.array(dtype=float),
-  body_parentid: wp.array(dtype=int),
-  body_rootid: wp.array(dtype=int),
-  body_weldid: wp.array(dtype=int),
-  body_dofnum: wp.array(dtype=int),
-  body_dofadr: wp.array(dtype=int),
-  body_invweight0: wp.array2d(dtype=wp.vec2),
-  dof_bodyid: wp.array(dtype=int),
-  dof_parentid: wp.array(dtype=int),
-  geom_bodyid: wp.array(dtype=int),
-  flex_vertadr: wp.array(dtype=int),
-  flex_vertbodyid: wp.array(dtype=int),
+  opt_impratio_invsqrt: wp.array[float],
+  body_parentid: wp.array[int],
+  body_rootid: wp.array[int],
+  body_weldid: wp.array[int],
+  body_dofnum: wp.array[int],
+  body_dofadr: wp.array[int],
+  body_invweight0: wp.array2d[wp.vec2],
+  dof_bodyid: wp.array[int],
+  dof_parentid: wp.array[int],
+  geom_bodyid: wp.array[int],
+  flex_vertadr: wp.array[int],
+  flex_vertbodyid: wp.array[int],
   is_sparse: bool,
   # Data in:
-  qvel_in: wp.array2d(dtype=float),
-  subtree_com_in: wp.array2d(dtype=wp.vec3),
-  cdof_in: wp.array2d(dtype=wp.spatial_vector),
+  qvel_in: wp.array2d[float],
+  subtree_com_in: wp.array2d[wp.vec3],
+  cdof_in: wp.array2d[wp.spatial_vector],
   njmax_in: int,
   njmax_nnz_in: int,
-  nacon_in: wp.array(dtype=int),
+  nacon_in: wp.array[int],
   # In:
-  dist_in: wp.array(dtype=float),
-  condim_in: wp.array(dtype=int),
-  includemargin_in: wp.array(dtype=float),
-  worldid_in: wp.array(dtype=int),
-  geom_in: wp.array(dtype=wp.vec2i),
-  flex_in: wp.array(dtype=wp.vec2i),
-  vert_in: wp.array(dtype=wp.vec2i),
-  pos_in: wp.array(dtype=wp.vec3),
-  frame_in: wp.array(dtype=wp.mat33),
-  friction_in: wp.array(dtype=vec5),
-  solref_in: wp.array(dtype=wp.vec2),
-  solreffriction_in: wp.array(dtype=wp.vec2),
-  solimp_in: wp.array(dtype=vec5),
-  type_in: wp.array(dtype=int),
+  dist_in: wp.array[float],
+  condim_in: wp.array[int],
+  includemargin_in: wp.array[float],
+  worldid_in: wp.array[int],
+  geom_in: wp.array[wp.vec2i],
+  flex_in: wp.array[wp.vec2i],
+  vert_in: wp.array[wp.vec2i],
+  pos_in: wp.array[wp.vec3],
+  frame_in: wp.array[wp.mat33],
+  friction_in: wp.array[vec5],
+  solref_in: wp.array[wp.vec2],
+  solreffriction_in: wp.array[wp.vec2],
+  solimp_in: wp.array[vec5],
+  type_in: wp.array[int],
   # Data out:
-  nefc_out: wp.array(dtype=int),
-  contact_efc_address_out: wp.array2d(dtype=int),
-  efc_type_out: wp.array2d(dtype=int),
-  efc_id_out: wp.array2d(dtype=int),
-  efc_J_rownnz_out: wp.array2d(dtype=int),
-  efc_J_rowadr_out: wp.array2d(dtype=int),
-  efc_J_colind_out: wp.array3d(dtype=int),
-  efc_J_out: wp.array3d(dtype=float),
-  efc_pos_out: wp.array2d(dtype=float),
-  efc_margin_out: wp.array2d(dtype=float),
-  efc_D_out: wp.array2d(dtype=float),
-  efc_vel_out: wp.array2d(dtype=float),
-  efc_aref_out: wp.array2d(dtype=float),
-  efc_frictionloss_out: wp.array2d(dtype=float),
+  nefc_out: wp.array[int],
+  contact_efc_address_out: wp.array2d[int],
+  efc_type_out: wp.array2d[int],
+  efc_id_out: wp.array2d[int],
+  efc_J_rownnz_out: wp.array2d[int],
+  efc_J_rowadr_out: wp.array2d[int],
+  efc_J_colind_out: wp.array3d[int],
+  efc_J_out: wp.array3d[float],
+  efc_pos_out: wp.array2d[float],
+  efc_margin_out: wp.array2d[float],
+  efc_D_out: wp.array2d[float],
+  efc_vel_out: wp.array2d[float],
+  efc_aref_out: wp.array2d[float],
+  efc_frictionloss_out: wp.array2d[float],
   # Out:
-  efc_nnz_out: wp.array(dtype=int),
+  efc_nnz_out: wp.array[int],
 ):
   conid, dimid = wp.tid()
 

mjx/mujoco/mjx/third_party/mujoco_warp/_src/derivative.py

@@ -31,27 +31,27 @@ wp.set_module_options({"enable_backward": False})
 @wp.kernel
 def _qderiv_actuator_passive_vel(
   # Model:
-  opt_timestep: wp.array(dtype=float),
-  actuator_dyntype: wp.array(dtype=int),
-  actuator_gaintype: wp.array(dtype=int),
-  actuator_biastype: wp.array(dtype=int),
-  actuator_actadr: wp.array(dtype=int),
-  actuator_actnum: wp.array(dtype=int),
-  actuator_forcelimited: wp.array(dtype=bool),
-  actuator_actlimited: wp.array(dtype=bool),
-  actuator_dynprm: wp.array2d(dtype=vec10f),
-  actuator_gainprm: wp.array2d(dtype=vec10f),
-  actuator_biasprm: wp.array2d(dtype=vec10f),
-  actuator_actearly: wp.array(dtype=bool),
-  actuator_forcerange: wp.array2d(dtype=wp.vec2),
-  actuator_actrange: wp.array2d(dtype=wp.vec2),
+  opt_timestep: wp.array[float],
+  actuator_dyntype: wp.array[int],
+  actuator_gaintype: wp.array[int],
+  actuator_biastype: wp.array[int],
+  actuator_actadr: wp.array[int],
+  actuator_actnum: wp.array[int],
+  actuator_forcelimited: wp.array[bool],
+  actuator_actlimited: wp.array[bool],
+  actuator_dynprm: wp.array2d[vec10f],
+  actuator_gainprm: wp.array2d[vec10f],
+  actuator_biasprm: wp.array2d[vec10f],
+  actuator_actearly: wp.array[bool],
+  actuator_forcerange: wp.array2d[wp.vec2],
+  actuator_actrange: wp.array2d[wp.vec2],
   # Data in:
-  act_in: wp.array2d(dtype=float),
-  ctrl_in: wp.array2d(dtype=float),
-  act_dot_in: wp.array2d(dtype=float),
-  actuator_force_in: wp.array2d(dtype=float),
+  act_in: wp.array2d[float],
+  ctrl_in: wp.array2d[float],
+  act_dot_in: wp.array2d[float],
+  actuator_force_in: wp.array2d[float],
   # Out:
-  vel_out: wp.array2d(dtype=float),
+  vel_out: wp.array2d[float],
 ):
   worldid, actid = wp.tid()
 
@@ -121,16 +121,16 @@ def _qderiv_actuator_passive_actuation_dense(
   # Model:
   nu: int,
   # Data in:
-  moment_rownnz_in: wp.array2d(dtype=int),
-  moment_rowadr_in: wp.array2d(dtype=int),
-  moment_colind_in: wp.array2d(dtype=int),
-  actuator_moment_in: wp.array2d(dtype=float),
+  moment_rownnz_in: wp.array2d[int],
+  moment_rowadr_in: wp.array2d[int],
+  moment_colind_in: wp.array2d[int],
+  actuator_moment_in: wp.array2d[float],
   # In:
-  vel_in: wp.array2d(dtype=float),
-  qMi: wp.array(dtype=int),
-  qMj: wp.array(dtype=int),
+  vel_in: wp.array2d[float],
+  qMi: wp.array[int],
+  qMj: wp.array[int],
   # Out:
-  qDeriv_out: wp.array3d(dtype=float),
+  qDeriv_out: wp.array3d[float],
 ):
   worldid, elemid = wp.tid()
 
@@ -171,18 +171,18 @@ def _qderiv_actuator_passive_actuation_dense(
 @wp.kernel
 def _qderiv_actuator_passive_actuation_sparse(
   # Model:
-  M_rownnz: wp.array(dtype=int),
-  M_rowadr: wp.array(dtype=int),
+  M_rownnz: wp.array[int],
+  M_rowadr: wp.array[int],
   # Data in:
-  moment_rownnz_in: wp.array2d(dtype=int),
-  moment_rowadr_in: wp.array2d(dtype=int),
-  moment_colind_in: wp.array2d(dtype=int),
-  actuator_moment_in: wp.array2d(dtype=float),
+  moment_rownnz_in: wp.array2d[int],
+  moment_rowadr_in: wp.array2d[int],
+  moment_colind_in: wp.array2d[int],
+  actuator_moment_in: wp.array2d[float],
   # In:
-  vel_in: wp.array2d(dtype=float),
-  qMj: wp.array(dtype=int),
+  vel_in: wp.array2d[float],
+  qMj: wp.array[int],
   # Out:
-  qDeriv_out: wp.array3d(dtype=float),
+  qDeriv_out: wp.array3d[float],
 ):
   worldid, actid = wp.tid()
 
@@ -225,18 +225,18 @@ def _qderiv_actuator_passive_actuation_sparse(
 @wp.kernel
 def _qderiv_actuator_passive(
   # Model:
-  opt_timestep: wp.array(dtype=float),
+  opt_timestep: wp.array[float],
   opt_disableflags: int,
-  dof_damping: wp.array2d(dtype=float),
+  dof_damping: wp.array2d[float],
   is_sparse: bool,
   # Data in:
-  qM_in: wp.array3d(dtype=float),
+  qM_in: wp.array3d[float],
   # In:
-  qMi: wp.array(dtype=int),
-  qMj: wp.array(dtype=int),
-  qDeriv_in: wp.array3d(dtype=float),
+  qMi: wp.array[int],
+  qMj: wp.array[int],
+  qDeriv_in: wp.array3d[float],
   # Out:
-  qDeriv_out: wp.array3d(dtype=float),
+  qDeriv_out: wp.array3d[float],
 ):
   worldid, elemid = wp.tid()
 
@@ -267,19 +267,19 @@ def _qderiv_actuator_passive(
 def _qderiv_tendon_damping(
   # Model:
   ntendon: int,
-  opt_timestep: wp.array(dtype=float),
-  ten_J_rownnz: wp.array(dtype=int),
-  ten_J_rowadr: wp.array(dtype=int),
-  ten_J_colind: wp.array(dtype=int),
-  tendon_damping: wp.array2d(dtype=float),
+  opt_timestep: wp.array[float],
+  ten_J_rownnz: wp.array[int],
+  ten_J_rowadr: wp.array[int],
+  ten_J_colind: wp.array[int],
+  tendon_damping: wp.array2d[float],
   is_sparse: bool,
   # Data in:
-  ten_J_in: wp.array2d(dtype=float),
+  ten_J_in: wp.array2d[float],
   # In:
-  qMi: wp.array(dtype=int),
-  qMj: wp.array(dtype=int),
+  qMi: wp.array[int],
+  qMj: wp.array[int],
   # Out:
-  qDeriv_out: wp.array3d(dtype=float),
+  qDeriv_out: wp.array3d[float],
 ):
   worldid, elemid = wp.tid()
   dofiid = qMi[elemid]
@@ -318,7 +318,7 @@ def _qderiv_tendon_damping(
 
 
 @event_scope
-def deriv_smooth_vel(m: Model, d: Data, out: wp.array2d(dtype=float)):
+def deriv_smooth_vel(m: Model, d: Data, out: wp.array2d[float]):
   """Analytical derivative of smooth forces w.r.t. velocities.
 
   Args:

mjx/mujoco/mjx/third_party/mujoco_warp/_src/forward.py

@@ -51,17 +51,17 @@ wp.set_module_options({"enable_backward": False})
 @wp.kernel
 def _next_position(
   # Model:
-  opt_timestep: wp.array(dtype=float),
-  jnt_type: wp.array(dtype=int),
-  jnt_qposadr: wp.array(dtype=int),
-  jnt_dofadr: wp.array(dtype=int),
+  opt_timestep: wp.array[float],
+  jnt_type: wp.array[int],
+  jnt_qposadr: wp.array[int],
+  jnt_dofadr: wp.array[int],
   # Data in:
-  qpos_in: wp.array2d(dtype=float),
-  qvel_in: wp.array2d(dtype=float),
+  qpos_in: wp.array2d[float],
+  qvel_in: wp.array2d[float],
   # In:
   qvel_scale_in: float,
   # Data out:
-  qpos_out: wp.array2d(dtype=float),
+  qpos_out: wp.array2d[float],
 ):
   worldid, jntid = wp.tid()
   timestep = opt_timestep[worldid % opt_timestep.shape[0]]
@@ -115,14 +115,14 @@ def _next_position(
 @wp.kernel
 def _next_velocity(
   # Model:
-  opt_timestep: wp.array(dtype=float),
+  opt_timestep: wp.array[float],
   # Data in:
-  qvel_in: wp.array2d(dtype=float),
-  qacc_in: wp.array2d(dtype=float),
+  qvel_in: wp.array2d[float],
+  qacc_in: wp.array2d[float],
   # In:
   qacc_scale_in: float,
   # Data out:
-  qvel_out: wp.array2d(dtype=float),
+  qvel_out: wp.array2d[float],
 ):
   worldid, dofid = wp.tid()
   timestep = opt_timestep[worldid % opt_timestep.shape[0]]
@@ -132,21 +132,21 @@ def _next_velocity(
 @wp.kernel
 def _next_activation(
   # Model:
-  opt_timestep: wp.array(dtype=float),
-  actuator_dyntype: wp.array(dtype=int),
-  actuator_actadr: wp.array(dtype=int),
-  actuator_actnum: wp.array(dtype=int),
-  actuator_actlimited: wp.array(dtype=bool),
-  actuator_dynprm: wp.array2d(dtype=vec10f),
-  actuator_actrange: wp.array2d(dtype=wp.vec2),
+  opt_timestep: wp.array[float],
+  actuator_dyntype: wp.array[int],
+  actuator_actadr: wp.array[int],
+  actuator_actnum: wp.array[int],
+  actuator_actlimited: wp.array[bool],
+  actuator_dynprm: wp.array2d[vec10f],
+  actuator_actrange: wp.array2d[wp.vec2],
   # Data in:
-  act_in: wp.array2d(dtype=float),
-  act_dot_in: wp.array2d(dtype=float),
+  act_in: wp.array2d[float],
+  act_dot_in: wp.array2d[float],
   # In:
   act_dot_scale: float,
   limit: bool,
   # Data out:
-  act_out: wp.array2d(dtype=float),
+  act_out: wp.array2d[float],
 ):
   worldid, uid = wp.tid()
   opt_timestep_id = worldid % opt_timestep.shape[0]
@@ -171,21 +171,21 @@ def _next_activation(
 @wp.kernel
 def _next_time(
   # Model:
-  opt_timestep: wp.array(dtype=float),
+  opt_timestep: wp.array[float],
   is_sparse: bool,
   # Data in:
-  nefc_in: wp.array(dtype=int),
-  time_in: wp.array(dtype=float),
-  efc_J_rownnz_in: wp.array2d(dtype=int),
-  efc_J_rowadr_in: wp.array2d(dtype=int),
+  nefc_in: wp.array[int],
+  time_in: wp.array[float],
+  efc_J_rownnz_in: wp.array2d[int],
+  efc_J_rowadr_in: wp.array2d[int],
   nworld_in: int,
   naconmax_in: int,
   njmax_in: int,
   njmax_nnz_in: int,
-  nacon_in: wp.array(dtype=int),
-  ncollision_in: wp.array(dtype=int),
+  nacon_in: wp.array[int],
+  ncollision_in: wp.array[int],
   # Data out:
-  time_out: wp.array(dtype=float),
+  time_out: wp.array[float],
 ):
   worldid = wp.tid()
   time_out[worldid] = time_in[worldid] + opt_timestep[worldid % opt_timestep.shape[0]]
@@ -277,11 +277,11 @@ def _advance(m: Model, d: Data, qacc: wp.array, qvel: Optional[wp.array] = None)
 @wp.kernel
 def _euler_damp_qfrc_sparse(
   # Model:
-  opt_timestep: wp.array(dtype=float),
-  dof_Madr: wp.array(dtype=int),
-  dof_damping: wp.array2d(dtype=float),
+  opt_timestep: wp.array[float],
+  dof_Madr: wp.array[int],
+  dof_damping: wp.array2d[float],
   # Out:
-  qM_integration_out: wp.array3d(dtype=float),
+  qM_integration_out: wp.array3d[float],
 ):
   worldid, tid = wp.tid()
   timestep = opt_timestep[worldid % opt_timestep.shape[0]]
@@ -295,15 +295,15 @@ def _tile_euler_dense(tile: TileSet):
   @wp.kernel(module="unique", enable_backward=False)
   def euler_dense(
     # Model:
-    opt_timestep: wp.array(dtype=float),
-    dof_damping: wp.array2d(dtype=float),
+    opt_timestep: wp.array[float],
+    dof_damping: wp.array2d[float],
     # Data in:
-    qM_in: wp.array3d(dtype=float),
-    efc_Ma_in: wp.array2d(dtype=float),
+    qM_in: wp.array3d[float],
+    efc_Ma_in: wp.array2d[float],
     # In:
-    adr_in: wp.array(dtype=int),
+    adr_in: wp.array[int],
     # Data out:
-    qacc_out: wp.array2d(dtype=float),
+    qacc_out: wp.array2d[float],
   ):
     worldid, nodeid = wp.tid()
     timestep = opt_timestep[worldid % opt_timestep.shape[0]]
@@ -358,8 +358,8 @@ def _rk_perturb_state(
   m: Model,
   d: Data,
   scale: float,
-  qpos_t0: wp.array2d(dtype=float),
-  qvel_t0: wp.array2d(dtype=float),
+  qpos_t0: wp.array2d[float],
+  qvel_t0: wp.array2d[float],
   act_t0: Optional[wp.array] = None,
 ):
   # position
@@ -403,13 +403,13 @@ def _rk_perturb_state(
 @wp.kernel
 def _rk_accumulate_velocity_acceleration(
   # Data in:
-  qvel_in: wp.array2d(dtype=float),
-  qacc_in: wp.array2d(dtype=float),
+  qvel_in: wp.array2d[float],
+  qacc_in: wp.array2d[float],
   # In:
   scale: float,
   # Data out:
-  qvel_out: wp.array2d(dtype=float),
-  qacc_out: wp.array2d(dtype=float),
+  qvel_out: wp.array2d[float],
+  qacc_out: wp.array2d[float],
 ):
   worldid, dofid = wp.tid()
   qvel_out[worldid, dofid] += scale * qvel_in[worldid, dofid]
@@ -419,11 +419,11 @@ def _rk_accumulate_velocity_acceleration(
 @wp.kernel
 def _rk_accumulate_activation_velocity(
   # Data in:
-  act_dot_in: wp.array2d(dtype=float),
+  act_dot_in: wp.array2d[float],
   # In:
   scale: float,
   # Data out:
-  act_dot_out: wp.array2d(dtype=float),
+  act_dot_out: wp.array2d[float],
 ):
   worldid, actid = wp.tid()
   act_dot_out[worldid, actid] += scale * act_dot_in[worldid, actid]
@@ -433,8 +433,8 @@ def _rk_accumulate(
   m: Model,
   d: Data,
   scale: float,
-  qvel_rk: wp.array2d(dtype=float),
-  qacc_rk: wp.array2d(dtype=float),
+  qvel_rk: wp.array2d[float],
+  qacc_rk: wp.array2d[float],
   act_dot_rk: Optional[wp.array] = None,
 ):
   """Computes one term of 1/6 k_1 + 1/3 k_2 + 1/3 k_3 + 1/6 k_4."""
@@ -540,13 +540,13 @@ def fwd_position(m: Model, d: Data, factorize: bool = True):
 @wp.kernel
 def _actuator_velocity(
   # Data in:
-  qvel_in: wp.array2d(dtype=float),
-  moment_rownnz_in: wp.array2d(dtype=int),
-  moment_rowadr_in: wp.array2d(dtype=int),
-  moment_colind_in: wp.array2d(dtype=int),
-  actuator_moment_in: wp.array2d(dtype=float),
+  qvel_in: wp.array2d[float],
+  moment_rownnz_in: wp.array2d[int],
+  moment_rowadr_in: wp.array2d[int],
+  moment_colind_in: wp.array2d[int],
+  actuator_moment_in: wp.array2d[float],
   # Data out:
-  actuator_velocity_out: wp.array2d(dtype=float),
+  actuator_velocity_out: wp.array2d[float],
 ):
   worldid, actid = wp.tid()
 
@@ -565,14 +565,14 @@ def _actuator_velocity(
 @wp.kernel
 def _tendon_velocity(
   # Model:
-  ten_J_rownnz: wp.array(dtype=int),
-  ten_J_rowadr: wp.array(dtype=int),
-  ten_J_colind: wp.array(dtype=int),
+  ten_J_rownnz: wp.array[int],
+  ten_J_rowadr: wp.array[int],
+  ten_J_colind: wp.array[int],
   # Data in:
-  qvel_in: wp.array2d(dtype=float),
-  ten_J_in: wp.array2d(dtype=float),
+  qvel_in: wp.array2d[float],
+  ten_J_in: wp.array2d[float],
   # Data out:
-  ten_velocity_out: wp.array2d(dtype=float),
+  ten_velocity_out: wp.array2d[float],
 ):
   worldid, tenid = wp.tid()
 
@@ -617,34 +617,34 @@ def fwd_velocity(m: Model, d: Data):
 def _actuator_force(
   # Model:
   na: int,
-  opt_timestep: wp.array(dtype=float),
-  actuator_dyntype: wp.array(dtype=int),
-  actuator_gaintype: wp.array(dtype=int),
-  actuator_biastype: wp.array(dtype=int),
-  actuator_actadr: wp.array(dtype=int),
-  actuator_actnum: wp.array(dtype=int),
-  actuator_ctrllimited: wp.array(dtype=bool),
-  actuator_forcelimited: wp.array(dtype=bool),
-  actuator_actlimited: wp.array(dtype=bool),
-  actuator_dynprm: wp.array2d(dtype=vec10f),
-  actuator_gainprm: wp.array2d(dtype=vec10f),
-  actuator_biasprm: wp.array2d(dtype=vec10f),
-  actuator_actearly: wp.array(dtype=bool),
-  actuator_ctrlrange: wp.array2d(dtype=wp.vec2),
-  actuator_forcerange: wp.array2d(dtype=wp.vec2),
-  actuator_actrange: wp.array2d(dtype=wp.vec2),
-  actuator_acc0: wp.array2d(dtype=float),
-  actuator_lengthrange: wp.array2d(dtype=wp.vec2),
+  opt_timestep: wp.array[float],
+  actuator_dyntype: wp.array[int],
+  actuator_gaintype: wp.array[int],
+  actuator_biastype: wp.array[int],
+  actuator_actadr: wp.array[int],
+  actuator_actnum: wp.array[int],
+  actuator_ctrllimited: wp.array[bool],
+  actuator_forcelimited: wp.array[bool],
+  actuator_actlimited: wp.array[bool],
+  actuator_dynprm: wp.array2d[vec10f],
+  actuator_gainprm: wp.array2d[vec10f],
+  actuator_biasprm: wp.array2d[vec10f],
+  actuator_actearly: wp.array[bool],
+  actuator_ctrlrange: wp.array2d[wp.vec2],
+  actuator_forcerange: wp.array2d[wp.vec2],
+  actuator_actrange: wp.array2d[wp.vec2],
+  actuator_acc0: wp.array2d[float],
+  actuator_lengthrange: wp.array2d[wp.vec2],
   # Data in:
-  act_in: wp.array2d(dtype=float),
-  ctrl_in: wp.array2d(dtype=float),
-  actuator_length_in: wp.array2d(dtype=float),
-  actuator_velocity_in: wp.array2d(dtype=float),
+  act_in: wp.array2d[float],
+  ctrl_in: wp.array2d[float],
+  actuator_length_in: wp.array2d[float],
+  actuator_velocity_in: wp.array2d[float],
   # In:
   dsbl_clampctrl: int,
   # Data out:
-  act_dot_out: wp.array2d(dtype=float),
-  actuator_force_out: wp.array2d(dtype=float),
+  act_dot_out: wp.array2d[float],
+  actuator_force_out: wp.array2d[float],
 ):
   worldid, uid = wp.tid()
 
@@ -738,12 +738,12 @@ def _actuator_force(
 @wp.kernel
 def _tendon_actuator_force(
   # Model:
-  actuator_trntype: wp.array(dtype=int),
-  actuator_trnid: wp.array(dtype=wp.vec2i),
+  actuator_trntype: wp.array[int],
+  actuator_trnid: wp.array[wp.vec2i],
   # Data in:
-  actuator_force_in: wp.array2d(dtype=float),
+  actuator_force_in: wp.array2d[float],
   # Out:
-  ten_actfrc_out: wp.array2d(dtype=float),
+  ten_actfrc_out: wp.array2d[float],
 ):
   worldid, actid = wp.tid()
 
@@ -756,14 +756,14 @@ def _tendon_actuator_force(
 @wp.kernel
 def _tendon_actuator_force_clamp(
   # Model:
-  tendon_actfrclimited: wp.array(dtype=bool),
-  tendon_actfrcrange: wp.array2d(dtype=wp.vec2),
-  actuator_trntype: wp.array(dtype=int),
-  actuator_trnid: wp.array(dtype=wp.vec2i),
+  tendon_actfrclimited: wp.array[bool],
+  tendon_actfrcrange: wp.array2d[wp.vec2],
+  actuator_trntype: wp.array[int],
+  actuator_trnid: wp.array[wp.vec2i],
   # In:
-  ten_actfrc_in: wp.array2d(dtype=float),
+  ten_actfrc_in: wp.array2d[float],
   # Data out:
-  actuator_force_out: wp.array2d(dtype=float),
+  actuator_force_out: wp.array2d[float],
 ):
   worldid, actid = wp.tid()
 
@@ -782,13 +782,13 @@ def _tendon_actuator_force_clamp(
 @wp.kernel
 def _qfrc_actuator(
   # Data in:
-  moment_rownnz_in: wp.array2d(dtype=int),
-  moment_rowadr_in: wp.array2d(dtype=int),
-  moment_colind_in: wp.array2d(dtype=int),
-  actuator_moment_in: wp.array2d(dtype=float),
-  actuator_force_in: wp.array2d(dtype=float),
+  moment_rownnz_in: wp.array2d[int],
+  moment_rowadr_in: wp.array2d[int],
+  moment_colind_in: wp.array2d[int],
+  actuator_moment_in: wp.array2d[float],
+  actuator_force_in: wp.array2d[float],
   # Data out:
-  qfrc_actuator_out: wp.array2d(dtype=float),
+  qfrc_actuator_out: wp.array2d[float],
 ):
   worldid, actid = wp.tid()
 
@@ -806,15 +806,15 @@ def _qfrc_actuator(
 def _qfrc_actuator_gravcomp_limits(
   # Model:
   ngravcomp: int,
-  jnt_actfrclimited: wp.array(dtype=bool),
-  jnt_actgravcomp: wp.array(dtype=int),
-  jnt_actfrcrange: wp.array2d(dtype=wp.vec2),
-  dof_jntid: wp.array(dtype=int),
+  jnt_actfrclimited: wp.array[bool],
+  jnt_actgravcomp: wp.array[int],
+  jnt_actfrcrange: wp.array2d[wp.vec2],
+  dof_jntid: wp.array[int],
   # Data in:
-  qfrc_gravcomp_in: wp.array2d(dtype=float),
-  qfrc_actuator_in: wp.array2d(dtype=float),
+  qfrc_gravcomp_in: wp.array2d[float],
+  qfrc_actuator_in: wp.array2d[float],
   # Data out:
-  qfrc_actuator_out: wp.array2d(dtype=float),
+  qfrc_actuator_out: wp.array2d[float],
 ):
   worldid, dofid = wp.tid()
   jntid = dof_jntid[dofid]
@@ -930,12 +930,12 @@ def fwd_actuation(m: Model, d: Data):
 @wp.kernel
 def _qfrc_smooth(
   # Data in:
-  qfrc_applied_in: wp.array2d(dtype=float),
-  qfrc_bias_in: wp.array2d(dtype=float),
-  qfrc_passive_in: wp.array2d(dtype=float),
-  qfrc_actuator_in: wp.array2d(dtype=float),
+  qfrc_applied_in: wp.array2d[float],
+  qfrc_bias_in: wp.array2d[float],
+  qfrc_passive_in: wp.array2d[float],
+  qfrc_actuator_in: wp.array2d[float],
   # Data out:
-  qfrc_smooth_out: wp.array2d(dtype=float),
+  qfrc_smooth_out: wp.array2d[float],
 ):
   worldid, dofid = wp.tid()
   qfrc_smooth_out[worldid, dofid] = (

mjx/mujoco/mjx/third_party/mujoco_warp/_src/inverse.py

@@ -34,12 +34,12 @@ wp.set_module_options({"enable_backward": False})
 @wp.kernel
 def _qfrc_eulerdamp(
   # Model:
-  opt_timestep: wp.array(dtype=float),
-  dof_damping: wp.array2d(dtype=float),
+  opt_timestep: wp.array[float],
+  dof_damping: wp.array2d[float],
   # Data in:
-  qacc_in: wp.array2d(dtype=float),
+  qacc_in: wp.array2d[float],
   # Out:
-  qfrc_out: wp.array2d(dtype=float),
+  qfrc_out: wp.array2d[float],
 ):
   worldid, dofid = wp.tid()
   timestep = opt_timestep[worldid % opt_timestep.shape[0]]
@@ -49,13 +49,13 @@ def _qfrc_eulerdamp(
 @wp.kernel
 def _qfrc_inverse(
   # Data in:
-  qfrc_bias_in: wp.array2d(dtype=float),
-  qfrc_passive_in: wp.array2d(dtype=float),
-  qfrc_constraint_in: wp.array2d(dtype=float),
+  qfrc_bias_in: wp.array2d[float],
+  qfrc_passive_in: wp.array2d[float],
+  qfrc_constraint_in: wp.array2d[float],
   # In:
-  Ma: wp.array2d(dtype=float),
+  Ma: wp.array2d[float],
   # Data out:
-  qfrc_inverse_out: wp.array2d(dtype=float),
+  qfrc_inverse_out: wp.array2d[float],
 ):
   worldid, dofid = wp.tid()
 
@@ -67,7 +67,7 @@ def _qfrc_inverse(
   qfrc_inverse_out[worldid, dofid] = qfrc_inverse
 
 
-def discrete_acc(m: Model, d: Data, qacc: wp.array2d(dtype=float)):
+def discrete_acc(m: Model, d: Data, qacc: wp.array2d[float]):
   """Convert discrete-time qacc to continuous-time qacc.
 
   Args:

mjx/mujoco/mjx/third_party/mujoco_warp/_src/io.py

@@ -34,14 +34,20 @@ from mujoco.mjx.third_party.mujoco_warp._src.types import vec10
 from mujoco.mjx.third_party.mujoco_warp._src.util_pkg import check_version
 
 
-def _create_array(data: Any, spec: wp.array, sizes: dict[str, int]) -> wp.array | None:
+def _is_array_spec(typ) -> bool:
+  """Check if a type annotation is an array spec (wp.array instance or bracket annotation)."""
+  return isinstance(typ, wp.array) or type(typ).__name__ == "_ArrayAnnotation"
+
+
+def _create_array(data: Any, spec, sizes: dict[str, int]) -> wp.array | None:
   """Creates a warp array and populates it with data.
 
   The array shape is determined by a field spec referencing MjModel / MjData array sizes.
   """
+  spec_shape = getattr(spec, "shape", (0,))
   shape = None
-  if spec.shape != (0,):
-    shape = tuple(sizes[dim] if isinstance(dim, str) else dim for dim in spec.shape)
+  if spec_shape != (0,):
+    shape = tuple(sizes[dim] if isinstance(dim, str) else dim for dim in spec_shape)
 
   if data is None and shape is None:
     return None  # nothing to do
@@ -50,7 +56,7 @@ def _create_array(data: Any, spec: wp.array, sizes: dict[str, int]) -> wp.array
   else:
     array = wp.array(np.array(data), dtype=spec.dtype, shape=shape)
 
-  if spec.shape[0] == "*":
+  if spec_shape and spec_shape[0] == "*":
     # add private attribute for JAX to determine which fields are batched
     array._is_batched = True
     # also set stride 0 to 0 which is expected legacy behavior (but is deprecated)
@@ -194,7 +200,7 @@ def put_model(mjm: mujoco.MjModel) -> types.Model:
 
   # place opt on device
   for f in dataclasses.fields(types.Option):
-    if isinstance(f.type, wp.array):
+    if _is_array_spec(f.type):
       setattr(opt, f.name, _create_array(getattr(opt, f.name), f.type, {"*": 1}))
     else:
       setattr(opt, f.name, f.type(getattr(opt, f.name)))
@@ -635,7 +641,7 @@ def put_model(mjm: mujoco.MjModel) -> types.Model:
   # place m on device
   sizes = dict({"*": 1}, **{f.name: getattr(m, f.name) for f in dataclasses.fields(types.Model) if f.type is int})
   for f in dataclasses.fields(types.Model):
-    if isinstance(f.type, wp.array):
+    if _is_array_spec(f.type):
       setattr(m, f.name, _create_array(getattr(m, f.name), f.type, sizes))
 
   return m
@@ -1459,7 +1465,7 @@ def reset_data(m: types.Model, d: types.Data, reset: Optional[wp.array] = None):
   """
 
   @wp.kernel(module="unique", enable_backward=False)
-  def reset_xfrc_applied(reset_in: wp.array(dtype=bool), xfrc_applied_out: wp.array2d(dtype=wp.spatial_vector)):
+  def reset_xfrc_applied(reset_in: wp.array[bool], xfrc_applied_out: wp.array2d[wp.spatial_vector]):
     worldid, bodyid, elemid = wp.tid()
 
     if wp.static(reset is not None):
@@ -1469,7 +1475,7 @@ def reset_data(m: types.Model, d: types.Data, reset: Optional[wp.array] = None):
     xfrc_applied_out[worldid, bodyid][elemid] = 0.0
 
   @wp.kernel(module="unique", enable_backward=False)
-  def reset_qM(reset_in: wp.array(dtype=bool), qM_out: wp.array3d(dtype=float)):
+  def reset_qM(reset_in: wp.array[bool], qM_out: wp.array3d[float]):
     worldid, elemid1, elemid2 = wp.tid()
 
     if wp.static(reset is not None):
@@ -1487,31 +1493,31 @@ def reset_data(m: types.Model, d: types.Data, reset: Optional[wp.array] = None):
     na: int,
     neq: int,
     nsensordata: int,
-    qpos0: wp.array2d(dtype=float),
-    eq_active0: wp.array(dtype=bool),
+    qpos0: wp.array2d[float],
+    eq_active0: wp.array[bool],
     # Data in:
     nworld_in: int,
     # In:
-    reset_in: wp.array(dtype=bool),
+    reset_in: wp.array[bool],
     # Data out:
-    solver_niter_out: wp.array(dtype=int),
-    ne_out: wp.array(dtype=int),
-    nf_out: wp.array(dtype=int),
-    nl_out: wp.array(dtype=int),
-    nefc_out: wp.array(dtype=int),
-    time_out: wp.array(dtype=float),
-    energy_out: wp.array(dtype=wp.vec2),
-    qpos_out: wp.array2d(dtype=float),
-    qvel_out: wp.array2d(dtype=float),
-    act_out: wp.array2d(dtype=float),
-    qacc_warmstart_out: wp.array2d(dtype=float),
-    ctrl_out: wp.array2d(dtype=float),
-    qfrc_applied_out: wp.array2d(dtype=float),
-    eq_active_out: wp.array2d(dtype=bool),
-    qacc_out: wp.array2d(dtype=float),
-    act_dot_out: wp.array2d(dtype=float),
-    sensordata_out: wp.array2d(dtype=float),
-    nacon_out: wp.array(dtype=int),
+    solver_niter_out: wp.array[int],
+    ne_out: wp.array[int],
+    nf_out: wp.array[int],
+    nl_out: wp.array[int],
+    nefc_out: wp.array[int],
+    time_out: wp.array[float],
+    energy_out: wp.array[wp.vec2],
+    qpos_out: wp.array2d[float],
+    qvel_out: wp.array2d[float],
+    act_out: wp.array2d[float],
+    qacc_warmstart_out: wp.array2d[float],
+    ctrl_out: wp.array2d[float],
+    qfrc_applied_out: wp.array2d[float],
+    eq_active_out: wp.array2d[bool],
+    qacc_out: wp.array2d[float],
+    act_dot_out: wp.array2d[float],
+    sensordata_out: wp.array2d[float],
+    nacon_out: wp.array[int],
   ):
     worldid = wp.tid()
 
@@ -1549,14 +1555,14 @@ def reset_data(m: types.Model, d: types.Data, reset: Optional[wp.array] = None):
   @wp.kernel(module="unique", enable_backward=False)
   def reset_mocap(
     # Model:
-    body_mocapid: wp.array(dtype=int),
-    body_pos: wp.array2d(dtype=wp.vec3),
-    body_quat: wp.array2d(dtype=wp.quat),
+    body_mocapid: wp.array[int],
+    body_pos: wp.array2d[wp.vec3],
+    body_quat: wp.array2d[wp.quat],
     # In:
-    reset_in: wp.array(dtype=bool),
+    reset_in: wp.array[bool],
     # Data out:
-    mocap_pos_out: wp.array2d(dtype=wp.vec3),
-    mocap_quat_out: wp.array2d(dtype=wp.quat),
+    mocap_pos_out: wp.array2d[wp.vec3],
+    mocap_quat_out: wp.array2d[wp.quat],
   ):
     worldid, bodyid = wp.tid()
 
@@ -1573,27 +1579,27 @@ def reset_data(m: types.Model, d: types.Data, reset: Optional[wp.array] = None):
   @wp.kernel(module="unique", enable_backward=False)
   def reset_contact(
     # Data in:
-    nacon_in: wp.array(dtype=int),
+    nacon_in: wp.array[int],
     # In:
-    reset_in: wp.array(dtype=bool),
+    reset_in: wp.array[bool],
     nefcaddress: int,
     # Data out:
-    contact_dist_out: wp.array(dtype=float),
-    contact_pos_out: wp.array(dtype=wp.vec3),
-    contact_frame_out: wp.array(dtype=wp.mat33),
-    contact_includemargin_out: wp.array(dtype=float),
-    contact_friction_out: wp.array(dtype=types.vec5),
-    contact_solref_out: wp.array(dtype=wp.vec2),
-    contact_solreffriction_out: wp.array(dtype=wp.vec2),
-    contact_solimp_out: wp.array(dtype=types.vec5),
-    contact_dim_out: wp.array(dtype=int),
-    contact_geom_out: wp.array(dtype=wp.vec2i),
-    contact_flex_out: wp.array(dtype=wp.vec2i),
-    contact_vert_out: wp.array(dtype=wp.vec2i),
-    contact_efc_address_out: wp.array2d(dtype=int),
-    contact_worldid_out: wp.array(dtype=int),
-    contact_type_out: wp.array(dtype=int),
-    contact_geomcollisionid_out: wp.array(dtype=int),
+    contact_dist_out: wp.array[float],
+    contact_pos_out: wp.array[wp.vec3],
+    contact_frame_out: wp.array[wp.mat33],
+    contact_includemargin_out: wp.array[float],
+    contact_friction_out: wp.array[types.vec5],
+    contact_solref_out: wp.array[wp.vec2],
+    contact_solreffriction_out: wp.array[wp.vec2],
+    contact_solimp_out: wp.array[types.vec5],
+    contact_dim_out: wp.array[int],
+    contact_geom_out: wp.array[wp.vec2i],
+    contact_flex_out: wp.array[wp.vec2i],
+    contact_vert_out: wp.array[wp.vec2i],
+    contact_efc_address_out: wp.array2d[int],
+    contact_worldid_out: wp.array[int],
+    contact_type_out: wp.array[int],
+    contact_geomcollisionid_out: wp.array[int],
   ):
     conid = wp.tid()
 
@@ -1697,8 +1703,8 @@ def reset_data(m: types.Model, d: types.Data, reset: Optional[wp.array] = None):
 # kernel_analyzer: off
 @wp.kernel
 def _init_subtreemass(
-  body_mass_in: wp.array2d(dtype=float),
-  body_subtreemass_out: wp.array2d(dtype=float),
+  body_mass_in: wp.array2d[float],
+  body_subtreemass_out: wp.array2d[float],
 ):
   worldid, bodyid = wp.tid()
   body_mass_id = worldid % body_mass_in.shape[0]
@@ -1708,9 +1714,9 @@ def _init_subtreemass(
 
 @wp.kernel
 def _accumulate_subtreemass(
-  body_parentid: wp.array(dtype=int),
-  body_subtreemass_io: wp.array2d(dtype=float),
-  body_tree_: wp.array(dtype=int),
+  body_parentid: wp.array[int],
+  body_subtreemass_io: wp.array2d[float],
+  body_tree_: wp.array[int],
 ):
   worldid, nodeid = wp.tid()
   body_subtreemass_id = worldid % body_subtreemass_io.shape[0]
@@ -1722,8 +1728,8 @@ def _accumulate_subtreemass(
 
 @wp.kernel
 def _copy_qpos0_to_qpos(
-  qpos0: wp.array2d(dtype=float),
-  qpos_out: wp.array2d(dtype=float),
+  qpos0: wp.array2d[float],
+  qpos_out: wp.array2d[float],
 ):
   worldid, i = wp.tid()
   qpos0_id = worldid % qpos0.shape[0]
@@ -1732,8 +1738,8 @@ def _copy_qpos0_to_qpos(
 
 @wp.kernel
 def _copy_tendon_length0(
-  ten_length_in: wp.array2d(dtype=float),
-  tendon_length0_out: wp.array2d(dtype=float),
+  ten_length_in: wp.array2d[float],
+  tendon_length0_out: wp.array2d[float],
 ):
   worldid, tenid = wp.tid()
   tendon_length0_id = worldid % tendon_length0_out.shape[0]
@@ -1744,9 +1750,9 @@ def _copy_tendon_length0(
 def _compute_meaninertia(
   nv: int,
   is_sparse: bool,
-  dof_Madr_in: wp.array(dtype=int),
-  qM_in: wp.array3d(dtype=float),
-  meaninertia_out: wp.array(dtype=float),
+  dof_Madr_in: wp.array[int],
+  qM_in: wp.array3d[float],
+  meaninertia_out: wp.array[float],
 ):
   """Compute mean diagonal inertia from qM at qpos0."""
   worldid = wp.tid()
@@ -1771,7 +1777,7 @@ def _compute_meaninertia(
 @wp.kernel
 def _set_unit_vector(
   dofid_target: int,
-  unit_vec_out: wp.array2d(dtype=float),
+  unit_vec_out: wp.array2d[float],
 ):
   worldid = wp.tid()
   nv = unit_vec_out.shape[1]
@@ -1785,8 +1791,8 @@ def _set_unit_vector(
 @wp.kernel
 def _extract_dof_A_diag(
   dofid: int,
-  result_vec_in: wp.array2d(dtype=float),
-  dof_A_diag_out: wp.array2d(dtype=float),
+  result_vec_in: wp.array2d[float],
+  dof_A_diag_out: wp.array2d[float],
 ):
   worldid = wp.tid()
   dof_A_diag_id = worldid % dof_A_diag_out.shape[0]
@@ -1795,11 +1801,11 @@ def _extract_dof_A_diag(
 
 @wp.kernel
 def _finalize_dof_invweight0(
-  dof_jntid: wp.array(dtype=int),
-  jnt_type: wp.array(dtype=int),
-  jnt_dofadr: wp.array(dtype=int),
-  dof_A_diag_in: wp.array2d(dtype=float),
-  dof_invweight0_out: wp.array2d(dtype=float),
+  dof_jntid: wp.array[int],
+  jnt_type: wp.array[int],
+  jnt_dofadr: wp.array[int],
+  dof_A_diag_in: wp.array2d[float],
+  dof_invweight0_out: wp.array2d[float],
 ):
   worldid, dofid = wp.tid()
   dof_invweight0_id = worldid % dof_invweight0_out.shape[0]
@@ -1842,15 +1848,15 @@ def _compute_body_jac_row(
   nv: int,
   bodyid_target: int,
   row_idx: int,
-  body_parentid: wp.array(dtype=int),
-  body_rootid: wp.array(dtype=int),
-  body_dofadr: wp.array(dtype=int),
-  body_dofnum: wp.array(dtype=int),
-  dof_parentid: wp.array(dtype=int),
-  subtree_com_in: wp.array2d(dtype=wp.vec3),
-  xipos_in: wp.array2d(dtype=wp.vec3),
-  cdof_in: wp.array2d(dtype=wp.spatial_vector),
-  body_jac_row_out: wp.array2d(dtype=float),
+  body_parentid: wp.array[int],
+  body_rootid: wp.array[int],
+  body_dofadr: wp.array[int],
+  body_dofnum: wp.array[int],
+  dof_parentid: wp.array[int],
+  subtree_com_in: wp.array2d[wp.vec3],
+  xipos_in: wp.array2d[wp.vec3],
+  cdof_in: wp.array2d[wp.spatial_vector],
+  body_jac_row_out: wp.array2d[float],
 ):
   worldid = wp.tid()
 
@@ -1901,9 +1907,9 @@ def _compute_body_A_diag_entry(
   nv: int,
   bodyid_target: int,
   row_idx: int,
-  body_jac_row_in: wp.array2d(dtype=float),
-  result_vec_in: wp.array2d(dtype=float),
-  body_A_diag_out: wp.array3d(dtype=float),
+  body_jac_row_in: wp.array2d[float],
+  result_vec_in: wp.array2d[float],
+  body_A_diag_out: wp.array3d[float],
 ):
   worldid = wp.tid()
   body_A_diag_id = worldid % body_A_diag_out.shape[0]
@@ -1916,9 +1922,9 @@ def _compute_body_A_diag_entry(
 
 @wp.kernel
 def _finalize_body_invweight0(
-  body_weldid: wp.array(dtype=int),
-  body_A_diag_in: wp.array3d(dtype=float),
-  body_invweight0_out: wp.array2d(dtype=wp.vec2),
+  body_weldid: wp.array[int],
+  body_A_diag_in: wp.array3d[float],
+  body_invweight0_out: wp.array2d[wp.vec2],
 ):
   worldid, bodyid = wp.tid()
   body_invweight0_id = worldid % body_invweight0_out.shape[0]
@@ -1953,11 +1959,11 @@ def _finalize_body_invweight0(
 @wp.kernel
 def _copy_tendon_jacobian(
   tenid_target: int,
-  ten_J_rownnz: wp.array(dtype=int),
-  ten_J_rowadr: wp.array(dtype=int),
-  ten_J_colind: wp.array(dtype=int),
-  ten_J_in: wp.array2d(dtype=float),
-  ten_J_vec_out: wp.array2d(dtype=float),
+  ten_J_rownnz: wp.array[int],
+  ten_J_rowadr: wp.array[int],
+  ten_J_colind: wp.array[int],
+  ten_J_in: wp.array2d[float],
+  ten_J_vec_out: wp.array2d[float],
 ):
   worldid = wp.tid()
   nv = ten_J_in.shape[2]
@@ -1971,15 +1977,15 @@ def _copy_tendon_jacobian(
 @wp.kernel
 def _compute_tendon_dot_product(
   # Model:
-  ten_J_rownnz: wp.array(dtype=int),
-  ten_J_rowadr: wp.array(dtype=int),
-  ten_J_colind: wp.array(dtype=int),
+  ten_J_rownnz: wp.array[int],
+  ten_J_rowadr: wp.array[int],
+  ten_J_colind: wp.array[int],
   # In:
   tenid_target: int,
-  ten_J_in: wp.array2d(dtype=float),
-  result_vec_in: wp.array2d(dtype=float),
+  ten_J_in: wp.array2d[float],
+  result_vec_in: wp.array2d[float],
   # Out:
-  tendon_invweight0_out: wp.array2d(dtype=float),
+  tendon_invweight0_out: wp.array2d[float],
 ):
   worldid = wp.tid()
   tendon_invweight0_id = worldid % tendon_invweight0_out.shape[0]
@@ -1997,15 +2003,15 @@ def _compute_tendon_dot_product(
 
 @wp.kernel
 def _compute_cam_pos0(
-  cam_bodyid: wp.array(dtype=int),
-  cam_targetbodyid: wp.array(dtype=int),
-  cam_xpos_in: wp.array2d(dtype=wp.vec3),
-  cam_xmat_in: wp.array2d(dtype=wp.mat33),
-  xpos_in: wp.array2d(dtype=wp.vec3),
-  subtree_com_in: wp.array2d(dtype=wp.vec3),
-  cam_pos0_out: wp.array2d(dtype=wp.vec3),
-  cam_poscom0_out: wp.array2d(dtype=wp.vec3),
-  cam_mat0_out: wp.array2d(dtype=wp.mat33),
+  cam_bodyid: wp.array[int],
+  cam_targetbodyid: wp.array[int],
+  cam_xpos_in: wp.array2d[wp.vec3],
+  cam_xmat_in: wp.array2d[wp.mat33],
+  xpos_in: wp.array2d[wp.vec3],
+  subtree_com_in: wp.array2d[wp.vec3],
+  cam_pos0_out: wp.array2d[wp.vec3],
+  cam_poscom0_out: wp.array2d[wp.vec3],
+  cam_mat0_out: wp.array2d[wp.mat33],
 ):
   worldid, camid = wp.tid()
   cam_pos0_id = worldid % cam_pos0_out.shape[0]
@@ -2023,15 +2029,15 @@ def _compute_cam_pos0(
 
 @wp.kernel
 def _compute_light_pos0(
-  light_bodyid: wp.array(dtype=int),
-  light_targetbodyid: wp.array(dtype=int),
-  light_xpos_in: wp.array2d(dtype=wp.vec3),
-  light_xdir_in: wp.array2d(dtype=wp.vec3),
-  xpos_in: wp.array2d(dtype=wp.vec3),
-  subtree_com_in: wp.array2d(dtype=wp.vec3),
-  light_pos0_out: wp.array2d(dtype=wp.vec3),
-  light_poscom0_out: wp.array2d(dtype=wp.vec3),
-  light_dir0_out: wp.array2d(dtype=wp.vec3),
+  light_bodyid: wp.array[int],
+  light_targetbodyid: wp.array[int],
+  light_xpos_in: wp.array2d[wp.vec3],
+  light_xdir_in: wp.array2d[wp.vec3],
+  xpos_in: wp.array2d[wp.vec3],
+  subtree_com_in: wp.array2d[wp.vec3],
+  light_pos0_out: wp.array2d[wp.vec3],
+  light_poscom0_out: wp.array2d[wp.vec3],
+  light_dir0_out: wp.array2d[wp.vec3],
 ):
   worldid, lightid = wp.tid()
   light_pos0_id = worldid % light_pos0_out.shape[0]
@@ -2050,11 +2056,11 @@ def _compute_light_pos0(
 @wp.kernel
 def _copy_actuator_moment(
   actid_target: int,
-  moment_rownnz_in: wp.array2d(dtype=int),
-  moment_rowadr_in: wp.array2d(dtype=int),
-  moment_colind_in: wp.array2d(dtype=int),
-  actuator_moment_in: wp.array2d(dtype=float),
-  act_moment_vec_out: wp.array2d(dtype=float),
+  moment_rownnz_in: wp.array2d[int],
+  moment_rowadr_in: wp.array2d[int],
+  moment_colind_in: wp.array2d[int],
+  actuator_moment_in: wp.array2d[float],
+  act_moment_vec_out: wp.array2d[float],
 ):
   worldid = wp.tid()
   nv = act_moment_vec_out.shape[1]
@@ -2072,8 +2078,8 @@ def _copy_actuator_moment(
 def _compute_actuator_acc0(
   actid_target: int,
   nv: int,
-  result_vec_in: wp.array2d(dtype=float),
-  actuator_acc0_out: wp.array2d(dtype=float),
+  result_vec_in: wp.array2d[float],
+  actuator_acc0_out: wp.array2d[float],
 ):
   worldid = wp.tid()
   norm_sq = float(0.0)
@@ -2084,11 +2090,11 @@ def _compute_actuator_acc0(
 
 @wp.kernel
 def _compute_dof_M0(
-  dof_bodyid: wp.array(dtype=int),
-  dof_armature: wp.array2d(dtype=float),
-  cdof_in: wp.array2d(dtype=wp.spatial_vector),
-  crb_in: wp.array2d(dtype=vec10),
-  dof_M0_out: wp.array2d(dtype=float),
+  dof_bodyid: wp.array[int],
+  dof_armature: wp.array2d[float],
+  cdof_in: wp.array2d[wp.spatial_vector],
+  crb_in: wp.array2d[vec10],
+  dof_M0_out: wp.array2d[float],
 ):
   worldid, dofid = wp.tid()
   bodyid = dof_bodyid[dofid]
@@ -2099,15 +2105,15 @@ def _compute_dof_M0(
 
 @wp.kernel
 def _resolve_dampratio(
-  actuator_biastype: wp.array(dtype=int),
-  actuator_gainprm: wp.array2d(dtype=types.vec10f),
-  moment_rownnz_in: wp.array2d(dtype=int),
-  moment_rowadr_in: wp.array2d(dtype=int),
-  moment_colind_in: wp.array2d(dtype=int),
-  actuator_moment_in: wp.array2d(dtype=float),
-  dof_M0_in: wp.array2d(dtype=float),
+  actuator_biastype: wp.array[int],
+  actuator_gainprm: wp.array2d[types.vec10f],
+  moment_rownnz_in: wp.array2d[int],
+  moment_rowadr_in: wp.array2d[int],
+  moment_colind_in: wp.array2d[int],
+  actuator_moment_in: wp.array2d[float],
+  dof_M0_in: wp.array2d[float],
   nv: int,
-  actuator_biasprm: wp.array2d(dtype=types.vec10f),
+  actuator_biasprm: wp.array2d[types.vec10f],
 ):
   worldid, actid = wp.tid()
   biastype = actuator_biastype[actid]
@@ -2150,15 +2156,15 @@ def _resolve_dampratio(
 
 @wp.kernel
 def _set_length_range(
-  actuator_trntype: wp.array(dtype=int),
-  actuator_trnid: wp.array(dtype=wp.vec2i),
-  actuator_gear: wp.array2d(dtype=wp.spatial_vector),
-  jnt_limited: wp.array(dtype=int),
-  jnt_range: wp.array2d(dtype=wp.vec2),
-  tendon_limited: wp.array(dtype=int),
-  tendon_range: wp.array2d(dtype=wp.vec2),
+  actuator_trntype: wp.array[int],
+  actuator_trnid: wp.array[wp.vec2i],
+  actuator_gear: wp.array2d[wp.spatial_vector],
+  jnt_limited: wp.array[int],
+  jnt_range: wp.array2d[wp.vec2],
+  tendon_limited: wp.array[int],
+  tendon_range: wp.array2d[wp.vec2],
   ntendon: int,
-  actuator_lengthrange_out: wp.array2d(dtype=wp.vec2),
+  actuator_lengthrange_out: wp.array2d[wp.vec2],
 ):
   worldid, actid = wp.tid()
   trntype = actuator_trntype[actid]
@@ -2515,7 +2521,13 @@ def override_model(model: types.Model | mujoco.MjModel, overrides: dict[str, Any
       "AUTO": mujoco.mjtJacobian.mjJAC_AUTO,
     },
   }
-  mjw_only_fields = {"opt.broadphase", "opt.broadphase_filter", "opt.ls_parallel", "opt.graph_conditional"}
+  mjw_only_fields = {
+    "opt.broadphase",
+    "opt.broadphase_filter",
+    "opt.ls_parallel",
+    "opt.graph_conditional",
+    "opt.contact_sensor_maxmatch",
+  }
   mj_only_fields = {"opt.jacobian"}
 
   if not isinstance(overrides, dict):
@@ -2626,7 +2638,7 @@ def _build_rays(
   intrinsic: wp.vec4,
   znear: float,
   # Out:
-  ray_out: wp.array(dtype=wp.vec3),
+  ray_out: wp.array[wp.vec3],
 ):
   xid, yid = wp.tid()
   ray_out[offset + xid + yid * img_w] = render_util.compute_ray(
@@ -2682,18 +2694,16 @@ def create_render_context(
     if callable(_cubql_avail) and _cubql_avail():
       constructor = "cubql"
 
-  # Mesh BVHs
+  # Mesh BVHs – build for all meshes so per-world variants are available
   nmesh = mjm.nmesh
   geom_enabled_mask = np.isin(mjm.geom_group, list(enabled_geom_groups))
-  mesh_geom_mask = geom_enabled_mask & (mjm.geom_type == types.GeomType.MESH) & (mjm.geom_dataid >= 0)
-  used_mesh_id = set(mjm.geom_dataid[mesh_geom_mask].astype(int))
   geom_enabled_idx = np.nonzero(geom_enabled_mask)[0]
 
   mesh_registry = {}
   mesh_bvh_id = [wp.uint64(0) for _ in range(nmesh)]
   mesh_bounds_size = [wp.vec3(0.0, 0.0, 0.0) for _ in range(nmesh)]
 
-  for mid in used_mesh_id:
+  for mid in range(nmesh):
     mesh, half = bvh.build_mesh_bvh(mjm, mid, constructor=constructor)
     mesh_registry[mesh.id] = mesh
     mesh_bvh_id[mid] = mesh.id

mjx/mujoco/mjx/third_party/mujoco_warp/_src/island.py

@@ -26,24 +26,24 @@ from mujoco.mjx.third_party.mujoco_warp._src.warp_util import event_scope
 def _tree_edges(
   # Model:
   nv: int,
-  body_treeid: wp.array(dtype=int),
-  jnt_dofadr: wp.array(dtype=int),
-  dof_treeid: wp.array(dtype=int),
-  geom_bodyid: wp.array(dtype=int),
-  site_bodyid: wp.array(dtype=int),
-  eq_type: wp.array(dtype=int),
-  eq_obj1id: wp.array(dtype=int),
-  eq_obj2id: wp.array(dtype=int),
-  eq_objtype: wp.array(dtype=int),
+  body_treeid: wp.array[int],
+  jnt_dofadr: wp.array[int],
+  dof_treeid: wp.array[int],
+  geom_bodyid: wp.array[int],
+  site_bodyid: wp.array[int],
+  eq_type: wp.array[int],
+  eq_obj1id: wp.array[int],
+  eq_obj2id: wp.array[int],
+  eq_objtype: wp.array[int],
   # Data in:
-  nefc_in: wp.array(dtype=int),
-  contact_geom_in: wp.array(dtype=wp.vec2i),
-  efc_type_in: wp.array2d(dtype=int),
-  efc_id_in: wp.array2d(dtype=int),
-  efc_J_in: wp.array3d(dtype=float),
+  nefc_in: wp.array[int],
+  contact_geom_in: wp.array[wp.vec2i],
+  efc_type_in: wp.array2d[int],
+  efc_id_in: wp.array2d[int],
+  efc_J_in: wp.array3d[float],
   njmax_in: int,
   # Out:
-  tree_tree: wp.array3d(dtype=int),  # kernel_analyzer: off
+  tree_tree: wp.array3d[int],  # kernel_analyzer: off
 ):
   """Find tree edges."""
   worldid, efcid = wp.tid()
@@ -151,7 +151,7 @@ def _tree_edges(
     wp.atomic_max(tree_tree, worldid, first_tree, first_tree, 1)
 
 
-def tree_edges(m: types.Model, d: types.Data, tree_tree: wp.array3d(dtype=int)):
+def tree_edges(m: types.Model, d: types.Data, tree_tree: wp.array3d[int]):
   """Compute tree-tree adjacency matrix."""
   tree_tree.zero_()
   wp.launch(
@@ -184,14 +184,14 @@ def _flood_fill(
   # Model:
   ntree: int,
   # In:
-  tree_tree_in: wp.array3d(dtype=int),
-  labels_in: wp.array2d(dtype=int),
-  stack_in: wp.array2d(dtype=int),
+  tree_tree_in: wp.array3d[int],
+  labels_in: wp.array2d[int],
+  stack_in: wp.array2d[int],
   # Data out:
-  nisland_out: wp.array(dtype=int),
-  tree_island_out: wp.array2d(dtype=int),
+  nisland_out: wp.array[int],
+  tree_island_out: wp.array2d[int],
   # Out:
-  stack_out: wp.array2d(dtype=int),
+  stack_out: wp.array2d[int],
 ):
   """DFS flood fill to discover islands using tree_tree matrix."""
   worldid = wp.tid()

mjx/mujoco/mjx/third_party/mujoco_warp/_src/passive.py

@@ -71,18 +71,18 @@ def _ellipsoid_max_moment(size: wp.vec3, dir: int) -> float:
 def _spring_damper_dof_passive(
   # Model:
   opt_disableflags: int,
-  qpos_spring: wp.array2d(dtype=float),
-  jnt_type: wp.array(dtype=int),
-  jnt_qposadr: wp.array(dtype=int),
-  jnt_dofadr: wp.array(dtype=int),
-  jnt_stiffness: wp.array2d(dtype=float),
-  dof_damping: wp.array2d(dtype=float),
+  qpos_spring: wp.array2d[float],
+  jnt_type: wp.array[int],
+  jnt_qposadr: wp.array[int],
+  jnt_dofadr: wp.array[int],
+  jnt_stiffness: wp.array2d[float],
+  dof_damping: wp.array2d[float],
   # Data in:
-  qpos_in: wp.array2d(dtype=float),
-  qvel_in: wp.array2d(dtype=float),
+  qpos_in: wp.array2d[float],
+  qvel_in: wp.array2d[float],
   # Data out:
-  qfrc_spring_out: wp.array2d(dtype=float),
-  qfrc_damper_out: wp.array2d(dtype=float),
+  qfrc_spring_out: wp.array2d[float],
+  qfrc_damper_out: wp.array2d[float],
 ):
   worldid, jntid = wp.tid()
   dofid = jnt_dofadr[jntid]
@@ -182,22 +182,22 @@ def _spring_damper_dof_passive(
 @wp.kernel
 def _spring_damper_tendon_passive(
   # Model:
-  ten_J_rownnz: wp.array(dtype=int),
-  ten_J_rowadr: wp.array(dtype=int),
-  ten_J_colind: wp.array(dtype=int),
-  tendon_stiffness: wp.array2d(dtype=float),
-  tendon_damping: wp.array2d(dtype=float),
-  tendon_lengthspring: wp.array2d(dtype=wp.vec2),
+  ten_J_rownnz: wp.array[int],
+  ten_J_rowadr: wp.array[int],
+  ten_J_colind: wp.array[int],
+  tendon_stiffness: wp.array2d[float],
+  tendon_damping: wp.array2d[float],
+  tendon_lengthspring: wp.array2d[wp.vec2],
   # Data in:
-  ten_J_in: wp.array2d(dtype=float),
-  ten_length_in: wp.array2d(dtype=float),
-  ten_velocity_in: wp.array2d(dtype=float),
+  ten_J_in: wp.array2d[float],
+  ten_length_in: wp.array2d[float],
+  ten_velocity_in: wp.array2d[float],
   # In:
   dsbl_spring: bool,
   dsbl_damper: bool,
   # Data out:
-  qfrc_spring_out: wp.array2d(dtype=float),
-  qfrc_damper_out: wp.array2d(dtype=float),
+  qfrc_spring_out: wp.array2d[float],
+  qfrc_damper_out: wp.array2d[float],
 ):
   worldid, tenid, dofid_sparse = wp.tid()
 
@@ -246,18 +246,18 @@ def _spring_damper_tendon_passive(
 @wp.kernel
 def _gravity_force(
   # Model:
-  opt_gravity: wp.array(dtype=wp.vec3),
-  body_parentid: wp.array(dtype=int),
-  body_rootid: wp.array(dtype=int),
-  body_mass: wp.array2d(dtype=float),
-  body_gravcomp: wp.array2d(dtype=float),
-  dof_bodyid: wp.array(dtype=int),
+  opt_gravity: wp.array[wp.vec3],
+  body_parentid: wp.array[int],
+  body_rootid: wp.array[int],
+  body_mass: wp.array2d[float],
+  body_gravcomp: wp.array2d[float],
+  dof_bodyid: wp.array[int],
   # Data in:
-  xipos_in: wp.array2d(dtype=wp.vec3),
-  subtree_com_in: wp.array2d(dtype=wp.vec3),
-  cdof_in: wp.array2d(dtype=wp.spatial_vector),
+  xipos_in: wp.array2d[wp.vec3],
+  subtree_com_in: wp.array2d[wp.vec3],
+  cdof_in: wp.array2d[wp.spatial_vector],
   # Data out:
-  qfrc_gravcomp_out: wp.array2d(dtype=float),
+  qfrc_gravcomp_out: wp.array2d[float],
 ):
   worldid, bodyid, dofid = wp.tid()
   bodyid += 1  # skip world body
@@ -275,27 +275,27 @@ def _gravity_force(
 @wp.kernel
 def _fluid_force(
   # Model:
-  opt_wind: wp.array(dtype=wp.vec3),
-  opt_density: wp.array(dtype=float),
-  opt_viscosity: wp.array(dtype=float),
-  body_rootid: wp.array(dtype=int),
-  body_geomnum: wp.array(dtype=int),
-  body_geomadr: wp.array(dtype=int),
-  body_mass: wp.array2d(dtype=float),
-  body_inertia: wp.array2d(dtype=wp.vec3),
-  geom_type: wp.array(dtype=int),
-  geom_size: wp.array2d(dtype=wp.vec3),
-  geom_fluid: wp.array2d(dtype=float),
-  body_fluid_ellipsoid: wp.array(dtype=bool),
+  opt_wind: wp.array[wp.vec3],
+  opt_density: wp.array[float],
+  opt_viscosity: wp.array[float],
+  body_rootid: wp.array[int],
+  body_geomnum: wp.array[int],
+  body_geomadr: wp.array[int],
+  body_mass: wp.array2d[float],
+  body_inertia: wp.array2d[wp.vec3],
+  geom_type: wp.array[int],
+  geom_size: wp.array2d[wp.vec3],
+  geom_fluid: wp.array2d[float],
+  body_fluid_ellipsoid: wp.array[bool],
   # Data in:
-  xipos_in: wp.array2d(dtype=wp.vec3),
-  ximat_in: wp.array2d(dtype=wp.mat33),
-  geom_xpos_in: wp.array2d(dtype=wp.vec3),
-  geom_xmat_in: wp.array2d(dtype=wp.mat33),
-  subtree_com_in: wp.array2d(dtype=wp.vec3),
-  cvel_in: wp.array2d(dtype=wp.spatial_vector),
+  xipos_in: wp.array2d[wp.vec3],
+  ximat_in: wp.array2d[wp.mat33],
+  geom_xpos_in: wp.array2d[wp.vec3],
+  geom_xmat_in: wp.array2d[wp.mat33],
+  subtree_com_in: wp.array2d[wp.vec3],
+  cvel_in: wp.array2d[wp.spatial_vector],
   # Out:
-  fluid_applied_out: wp.array2d(dtype=wp.spatial_vector),
+  fluid_applied_out: wp.array2d[wp.spatial_vector],
 ):
   """Computes body-space fluid forces for both inertia-box and ellipsoid models."""
   worldid, bodyid = wp.tid()
@@ -535,18 +535,18 @@ def _fluid(m: Model, d: Data):
 @wp.kernel
 def _qfrc_passive(
   # Model:
-  jnt_actgravcomp: wp.array(dtype=int),
-  dof_jntid: wp.array(dtype=int),
+  jnt_actgravcomp: wp.array[int],
+  dof_jntid: wp.array[int],
   has_fluid: bool,
   # Data in:
-  qfrc_spring_in: wp.array2d(dtype=float),
-  qfrc_damper_in: wp.array2d(dtype=float),
-  qfrc_gravcomp_in: wp.array2d(dtype=float),
-  qfrc_fluid_in: wp.array2d(dtype=float),
+  qfrc_spring_in: wp.array2d[float],
+  qfrc_damper_in: wp.array2d[float],
+  qfrc_gravcomp_in: wp.array2d[float],
+  qfrc_fluid_in: wp.array2d[float],
   # In:
   gravcomp: bool,
   # Data out:
-  qfrc_passive_out: wp.array2d(dtype=float),
+  qfrc_passive_out: wp.array2d[float],
 ):
   worldid, dofid = wp.tid()
   qfrc_passive = qfrc_spring_in[worldid, dofid]
@@ -567,29 +567,29 @@ def _qfrc_passive(
 def _flex_elasticity(
   # Model:
   nflex: int,
-  opt_timestep: wp.array(dtype=float),
-  body_dofadr: wp.array(dtype=int),
-  flex_dim: wp.array(dtype=int),
-  flex_vertadr: wp.array(dtype=int),
-  flex_edgeadr: wp.array(dtype=int),
-  flex_elemadr: wp.array(dtype=int),
-  flex_elemnum: wp.array(dtype=int),
-  flex_elemdataadr: wp.array(dtype=int),
-  flex_elemedgeadr: wp.array(dtype=int),
-  flex_vertbodyid: wp.array(dtype=int),
-  flex_elem: wp.array(dtype=int),
-  flex_elemedge: wp.array(dtype=int),
-  flexedge_length0: wp.array(dtype=float),
-  flex_stiffness: wp.array2d(dtype=float),
-  flex_damping: wp.array(dtype=float),
+  opt_timestep: wp.array[float],
+  body_dofadr: wp.array[int],
+  flex_dim: wp.array[int],
+  flex_vertadr: wp.array[int],
+  flex_edgeadr: wp.array[int],
+  flex_elemadr: wp.array[int],
+  flex_elemnum: wp.array[int],
+  flex_elemdataadr: wp.array[int],
+  flex_elemedgeadr: wp.array[int],
+  flex_vertbodyid: wp.array[int],
+  flex_elem: wp.array[int],
+  flex_elemedge: wp.array[int],
+  flexedge_length0: wp.array[float],
+  flex_stiffness: wp.array2d[float],
+  flex_damping: wp.array[float],
   # Data in:
-  flexvert_xpos_in: wp.array2d(dtype=wp.vec3),
-  flexedge_length_in: wp.array2d(dtype=float),
-  flexedge_velocity_in: wp.array2d(dtype=float),
+  flexvert_xpos_in: wp.array2d[wp.vec3],
+  flexedge_length_in: wp.array2d[float],
+  flexedge_velocity_in: wp.array2d[float],
   # In:
   dsbl_damper: bool,
   # Data out:
-  qfrc_spring_out: wp.array2d(dtype=float),
+  qfrc_spring_out: wp.array2d[float],
 ):
   worldid, elemid = wp.tid()
   timestep = opt_timestep[worldid % opt_timestep.shape[0]]
@@ -665,19 +665,19 @@ def _flex_elasticity(
 def _flex_bending(
   # Model:
   nflex: int,
-  body_dofadr: wp.array(dtype=int),
-  flex_dim: wp.array(dtype=int),
-  flex_vertadr: wp.array(dtype=int),
-  flex_edgeadr: wp.array(dtype=int),
-  flex_edgenum: wp.array(dtype=int),
-  flex_vertbodyid: wp.array(dtype=int),
-  flex_edge: wp.array(dtype=wp.vec2i),
-  flex_edgeflap: wp.array(dtype=wp.vec2i),
-  flex_bending: wp.array2d(dtype=float),
+  body_dofadr: wp.array[int],
+  flex_dim: wp.array[int],
+  flex_vertadr: wp.array[int],
+  flex_edgeadr: wp.array[int],
+  flex_edgenum: wp.array[int],
+  flex_vertbodyid: wp.array[int],
+  flex_edge: wp.array[wp.vec2i],
+  flex_edgeflap: wp.array[wp.vec2i],
+  flex_bending: wp.array2d[float],
   # Data in:
-  flexvert_xpos_in: wp.array2d(dtype=wp.vec3),
+  flexvert_xpos_in: wp.array2d[wp.vec3],
   # Data out:
-  qfrc_spring_out: wp.array2d(dtype=float),
+  qfrc_spring_out: wp.array2d[float],
 ):
   worldid, edgeid = wp.tid()
   nvert = 4

mjx/mujoco/mjx/third_party/mujoco_warp/_src/ray.py

@@ -52,12 +52,12 @@ def _ray_map(pos: wp.vec3, mat: wp.mat33, pnt: wp.vec3, vec: wp.vec3) -> Tuple[w
 @wp.func
 def _ray_eliminate(
   # Model:
-  body_weldid: wp.array(dtype=int),
-  geom_bodyid: wp.array(dtype=int),
-  geom_matid: wp.array(dtype=int),  # kernel_analyzer: ignore
-  geom_group: wp.array(dtype=int),
-  geom_rgba: wp.array(dtype=wp.vec4),  # kernel_analyzer: ignore
-  mat_rgba: wp.array(dtype=wp.vec4),  # kernel_analyzer: ignore
+  body_weldid: wp.array[int],
+  geom_bodyid: wp.array[int],
+  geom_matid: wp.array[int],  # kernel_analyzer: ignore
+  geom_group: wp.array[int],
+  geom_rgba: wp.array[wp.vec4],  # kernel_analyzer: ignore
+  mat_rgba: wp.array[wp.vec4],  # kernel_analyzer: ignore
   # In:
   geomid: int,
   geomgroup: vec6,
@@ -451,26 +451,27 @@ def ray_box(pos: wp.vec3, mat: wp.mat33, size: wp.vec3, pnt: wp.vec3, vec: wp.ve
 @wp.func
 def ray_hfield(
   # Model:
-  geom_type: wp.array(dtype=int),
-  geom_dataid: wp.array(dtype=int),
-  hfield_size: wp.array(dtype=wp.vec4),
-  hfield_nrow: wp.array(dtype=int),
-  hfield_ncol: wp.array(dtype=int),
-  hfield_adr: wp.array(dtype=int),
-  hfield_data: wp.array(dtype=float),
+  geom_type: wp.array[int],
+  geom_dataid: wp.array2d[int],
+  hfield_size: wp.array[wp.vec4],
+  hfield_nrow: wp.array[int],
+  hfield_ncol: wp.array[int],
+  hfield_adr: wp.array[int],
+  hfield_data: wp.array[float],
   # In:
   pos: wp.vec3,
   mat: wp.mat33,
   pnt: wp.vec3,
   vec: wp.vec3,
   id: int,
+  worldid: int,
 ) -> Tuple[float, wp.vec3]:
   # check geom type
   if geom_type[id] != GeomType.HFIELD:
     return -1.0, wp.vec3()
 
   # hfield id and dimensions
-  hid = geom_dataid[id]
+  hid = geom_dataid[worldid % geom_dataid.shape[0], id]
   nrow = hfield_nrow[hid]
   ncol = hfield_ncol[hid]
 
@@ -622,10 +623,10 @@ def ray_hfield(
 def ray_mesh(
   # Model:
   nmeshface: int,
-  mesh_vertadr: wp.array(dtype=int),
-  mesh_faceadr: wp.array(dtype=int),
-  mesh_vert: wp.array(dtype=wp.vec3),
-  mesh_face: wp.array(dtype=wp.vec3i),
+  mesh_vertadr: wp.array[int],
+  mesh_faceadr: wp.array[int],
+  mesh_vert: wp.array[wp.vec3],
+  mesh_face: wp.array[wp.vec3i],
   # In:
   data_id: int,
   pos: wp.vec3,
@@ -699,7 +700,7 @@ def ray_mesh(
 @wp.func
 def ray_mesh_with_bvh(
   # In:
-  mesh_bvh_id: wp.array(dtype=wp.uint64),
+  mesh_bvh_id: wp.array[wp.uint64],
   mesh_geom_id: int,
   pos: wp.vec3,
   mat: wp.mat33,
@@ -732,7 +733,7 @@ def ray_mesh_with_bvh(
 @wp.func
 def ray_mesh_with_bvh_anyhit(
   # In:
-  mesh_bvh_id: wp.array(dtype=wp.uint64),
+  mesh_bvh_id: wp.array[wp.uint64],
   mesh_geom_id: int,
   pos: wp.vec3,
   mat: wp.mat33,
@@ -752,7 +753,7 @@ def ray_mesh_with_bvh_anyhit(
 @wp.func
 def ray_flex_with_bvh(
   # In:
-  flex_bvh_id: wp.array(dtype=wp.uint64),
+  flex_bvh_id: wp.array[wp.uint64],
   flexid: int,
   group_root: int,
   pnt: wp.vec3,
@@ -781,7 +782,7 @@ def ray_flex_with_bvh(
 @wp.func
 def ray_flex_with_bvh_anyhit(
   # In:
-  flex_bvh_id: wp.array(dtype=wp.uint64),
+  flex_bvh_id: wp.array[wp.uint64],
   flexid: int,
   group_root: int,
   pnt: wp.vec3,
@@ -823,27 +824,27 @@ def ray_geom(pos: wp.vec3, mat: wp.mat33, size: wp.vec3, pnt: wp.vec3, vec: wp.v
 def _ray_geom_mesh(
   # Model:
   nmeshface: int,
-  body_weldid: wp.array(dtype=int),
-  geom_type: wp.array(dtype=int),
-  geom_bodyid: wp.array(dtype=int),
-  geom_dataid: wp.array(dtype=int),
-  geom_matid: wp.array2d(dtype=int),
-  geom_group: wp.array(dtype=int),
-  geom_size: wp.array2d(dtype=wp.vec3),
-  geom_rgba: wp.array2d(dtype=wp.vec4),
-  mesh_vertadr: wp.array(dtype=int),
-  mesh_faceadr: wp.array(dtype=int),
-  mesh_vert: wp.array(dtype=wp.vec3),
-  mesh_face: wp.array(dtype=wp.vec3i),
-  hfield_size: wp.array(dtype=wp.vec4),
-  hfield_nrow: wp.array(dtype=int),
-  hfield_ncol: wp.array(dtype=int),
-  hfield_adr: wp.array(dtype=int),
-  hfield_data: wp.array(dtype=float),
-  mat_rgba: wp.array2d(dtype=wp.vec4),
+  body_weldid: wp.array[int],
+  geom_type: wp.array[int],
+  geom_bodyid: wp.array[int],
+  geom_dataid: wp.array2d[int],
+  geom_matid: wp.array2d[int],
+  geom_group: wp.array[int],
+  geom_size: wp.array2d[wp.vec3],
+  geom_rgba: wp.array2d[wp.vec4],
+  mesh_vertadr: wp.array[int],
+  mesh_faceadr: wp.array[int],
+  mesh_vert: wp.array[wp.vec3],
+  mesh_face: wp.array[wp.vec3i],
+  hfield_size: wp.array[wp.vec4],
+  hfield_nrow: wp.array[int],
+  hfield_ncol: wp.array[int],
+  hfield_adr: wp.array[int],
+  hfield_data: wp.array[float],
+  mat_rgba: wp.array2d[wp.vec4],
   # Data in:
-  geom_xpos_in: wp.array2d(dtype=wp.vec3),
-  geom_xmat_in: wp.array2d(dtype=wp.mat33),
+  geom_xpos_in: wp.array2d[wp.vec3],
+  geom_xmat_in: wp.array2d[wp.mat33],
   # In:
   worldid: int,
   pnt: wp.vec3,
@@ -876,7 +877,7 @@ def _ray_geom_mesh(
         mesh_faceadr,
         mesh_vert,
         mesh_face,
-        geom_dataid[geomid],
+        geom_dataid[worldid % geom_dataid.shape[0], geomid],
         pos,
         mat,
         geom_size[worldid % geom_size.shape[0], geomid],
@@ -897,6 +898,7 @@ def _ray_geom_mesh(
         pnt,
         vec,
         geomid,
+        worldid,
       )
     else:
       return ray_geom(pos, mat, geom_size[worldid % geom_size.shape[0], geomid], pnt, vec, type)
@@ -909,37 +911,37 @@ def _ray(
   # Model:
   ngeom: int,
   nmeshface: int,
-  body_weldid: wp.array(dtype=int),
-  geom_type: wp.array(dtype=int),
-  geom_bodyid: wp.array(dtype=int),
-  geom_dataid: wp.array(dtype=int),
-  geom_matid: wp.array2d(dtype=int),
-  geom_group: wp.array(dtype=int),
-  geom_size: wp.array2d(dtype=wp.vec3),
-  geom_rgba: wp.array2d(dtype=wp.vec4),
-  mesh_vertadr: wp.array(dtype=int),
-  mesh_faceadr: wp.array(dtype=int),
-  mesh_vert: wp.array(dtype=wp.vec3),
-  mesh_face: wp.array(dtype=wp.vec3i),
-  hfield_size: wp.array(dtype=wp.vec4),
-  hfield_nrow: wp.array(dtype=int),
-  hfield_ncol: wp.array(dtype=int),
-  hfield_adr: wp.array(dtype=int),
-  hfield_data: wp.array(dtype=float),
-  mat_rgba: wp.array2d(dtype=wp.vec4),
+  body_weldid: wp.array[int],
+  geom_type: wp.array[int],
+  geom_bodyid: wp.array[int],
+  geom_dataid: wp.array2d[int],
+  geom_matid: wp.array2d[int],
+  geom_group: wp.array[int],
+  geom_size: wp.array2d[wp.vec3],
+  geom_rgba: wp.array2d[wp.vec4],
+  mesh_vertadr: wp.array[int],
+  mesh_faceadr: wp.array[int],
+  mesh_vert: wp.array[wp.vec3],
+  mesh_face: wp.array[wp.vec3i],
+  hfield_size: wp.array[wp.vec4],
+  hfield_nrow: wp.array[int],
+  hfield_ncol: wp.array[int],
+  hfield_adr: wp.array[int],
+  hfield_data: wp.array[float],
+  mat_rgba: wp.array2d[wp.vec4],
   # Data in:
-  geom_xpos_in: wp.array2d(dtype=wp.vec3),
-  geom_xmat_in: wp.array2d(dtype=wp.mat33),
+  geom_xpos_in: wp.array2d[wp.vec3],
+  geom_xmat_in: wp.array2d[wp.mat33],
   # In:
-  pnt: wp.array2d(dtype=wp.vec3),
-  vec: wp.array2d(dtype=wp.vec3),
+  pnt: wp.array2d[wp.vec3],
+  vec: wp.array2d[wp.vec3],
   geomgroup: vec6,
   flg_static: bool,
-  bodyexclude: wp.array(dtype=int),
+  bodyexclude: wp.array[int],
   # Out:
-  dist_out: wp.array2d(dtype=float),
-  geomid_out: wp.array2d(dtype=int),
-  normal_out: wp.array2d(dtype=wp.vec3),
+  dist_out: wp.array2d[float],
+  geomid_out: wp.array2d[int],
+  normal_out: wp.array2d[wp.vec3],
 ):
   worldid, rayid, tid = wp.tid()
 
@@ -1011,18 +1013,18 @@ def _ray(
 @wp.func
 def _ray_geom_mesh_bvh(
   # Model:
-  body_weldid: wp.array(dtype=int),
-  geom_type: wp.array(dtype=int),
-  geom_bodyid: wp.array(dtype=int),
-  geom_dataid: wp.array(dtype=int),
-  geom_matid: wp.array2d(dtype=int),
-  geom_group: wp.array(dtype=int),
-  geom_size: wp.array2d(dtype=wp.vec3),
-  geom_rgba: wp.array2d(dtype=wp.vec4),
-  mat_rgba: wp.array2d(dtype=wp.vec4),
+  body_weldid: wp.array[int],
+  geom_type: wp.array[int],
+  geom_bodyid: wp.array[int],
+  geom_dataid: wp.array2d[int],
+  geom_matid: wp.array2d[int],
+  geom_group: wp.array[int],
+  geom_size: wp.array2d[wp.vec3],
+  geom_rgba: wp.array2d[wp.vec4],
+  mat_rgba: wp.array2d[wp.vec4],
   # Data in:
-  geom_xpos_in: wp.array2d(dtype=wp.vec3),
-  geom_xmat_in: wp.array2d(dtype=wp.mat33),
+  geom_xpos_in: wp.array2d[wp.vec3],
+  geom_xmat_in: wp.array2d[wp.mat33],
   # In:
   worldid: int,
   pnt: wp.vec3,
@@ -1031,8 +1033,8 @@ def _ray_geom_mesh_bvh(
   flg_static: bool,
   bodyexclude: int,
   geomid: int,
-  mesh_bvh_id: wp.array(dtype=wp.uint64),
-  hfield_bvh_id: wp.array(dtype=wp.uint64),
+  mesh_bvh_id: wp.array[wp.uint64],
+  hfield_bvh_id: wp.array[wp.uint64],
   min_dist: float,
 ) -> Tuple[float, wp.vec3]:
   if not _ray_eliminate(
@@ -1055,7 +1057,7 @@ def _ray_geom_mesh_bvh(
       bvh_ids = mesh_bvh_id if gtype == GeomType.MESH else hfield_bvh_id
       t, n, u, v, f, geom_mesh_id = ray_mesh_with_bvh(
         bvh_ids,
-        geom_dataid[geomid],
+        geom_dataid[worldid % geom_dataid.shape[0], geomid],
         pos,
         mat,
         pnt,
@@ -1081,33 +1083,33 @@ def _ray_geom_mesh_bvh(
 def _ray_bvh(
   # Model:
   ngeom: int,
-  body_weldid: wp.array(dtype=int),
-  geom_type: wp.array(dtype=int),
-  geom_bodyid: wp.array(dtype=int),
-  geom_dataid: wp.array(dtype=int),
-  geom_matid: wp.array2d(dtype=int),
-  geom_group: wp.array(dtype=int),
-  geom_size: wp.array2d(dtype=wp.vec3),
-  geom_rgba: wp.array2d(dtype=wp.vec4),
-  mat_rgba: wp.array2d(dtype=wp.vec4),
+  body_weldid: wp.array[int],
+  geom_type: wp.array[int],
+  geom_bodyid: wp.array[int],
+  geom_dataid: wp.array2d[int],
+  geom_matid: wp.array2d[int],
+  geom_group: wp.array[int],
+  geom_size: wp.array2d[wp.vec3],
+  geom_rgba: wp.array2d[wp.vec4],
+  mat_rgba: wp.array2d[wp.vec4],
   # Data in:
-  geom_xpos_in: wp.array2d(dtype=wp.vec3),
-  geom_xmat_in: wp.array2d(dtype=wp.mat33),
+  geom_xpos_in: wp.array2d[wp.vec3],
+  geom_xmat_in: wp.array2d[wp.mat33],
   # In:
-  pnt: wp.array2d(dtype=wp.vec3),
-  vec: wp.array2d(dtype=wp.vec3),
+  pnt: wp.array2d[wp.vec3],
+  vec: wp.array2d[wp.vec3],
   geomgroup: vec6,
   flg_static: bool,
-  bodyexclude: wp.array(dtype=int),
+  bodyexclude: wp.array[int],
   bvh_id: wp.uint64,
-  group_root: wp.array(dtype=int),
-  enabled_geom_ids: wp.array(dtype=int),
-  mesh_bvh_id: wp.array(dtype=wp.uint64),
-  hfield_bvh_id: wp.array(dtype=wp.uint64),
+  group_root: wp.array[int],
+  enabled_geom_ids: wp.array[int],
+  mesh_bvh_id: wp.array[wp.uint64],
+  hfield_bvh_id: wp.array[wp.uint64],
   # Out:
-  dist_out: wp.array2d(dtype=float),
-  geomid_out: wp.array2d(dtype=int),
-  normal_out: wp.array2d(dtype=wp.vec3),
+  dist_out: wp.array2d[float],
+  geomid_out: wp.array2d[int],
+  normal_out: wp.array2d[wp.vec3],
 ):
   worldid, rayid = wp.tid()
 
@@ -1166,8 +1168,8 @@ def _ray_bvh(
 def ray(
   m: Model,
   d: Data,
-  pnt: wp.array2d(dtype=wp.vec3),
-  vec: wp.array2d(dtype=wp.vec3),
+  pnt: wp.array2d[wp.vec3],
+  vec: wp.array2d[wp.vec3],
   geomgroup: vec6 | None = None,
   flg_static: bool = True,
   bodyexclude: int = -1,
@@ -1210,14 +1212,14 @@ def ray(
 def rays(
   m: Model,
   d: Data,
-  pnt: wp.array2d(dtype=wp.vec3),
-  vec: wp.array2d(dtype=wp.vec3),
+  pnt: wp.array2d[wp.vec3],
+  vec: wp.array2d[wp.vec3],
   geomgroup: vec6,
   flg_static: bool,
-  bodyexclude: wp.array(dtype=int),
-  dist: wp.array2d(dtype=float),
-  geomid: wp.array2d(dtype=int),
-  normal: wp.array2d(dtype=wp.vec3),
+  bodyexclude: wp.array[int],
+  dist: wp.array2d[float],
+  geomid: wp.array2d[int],
+  normal: wp.array2d[wp.vec3],
   rc: RenderContext | None = None,
 ):
   """Ray intersection for multiple worlds and multiple rays.

mjx/mujoco/mjx/third_party/mujoco_warp/_src/render.py

@@ -43,17 +43,17 @@ wp.set_module_options({"enable_backward": False})
 @wp.func
 def sample_texture(
   # Model:
-  geom_type: wp.array(dtype=int),
-  mesh_faceadr: wp.array(dtype=int),
+  geom_type: wp.array[int],
+  mesh_faceadr: wp.array[int],
   # In:
   geom_id: int,
   tex_repeat: wp.vec2,
   tex: wp.Texture2D,
   pos: wp.vec3,
   rot: wp.mat33,
-  mesh_facetexcoord: wp.array(dtype=wp.vec3i),
-  mesh_texcoord: wp.array(dtype=wp.vec2),
-  mesh_texcoord_offsets: wp.array(dtype=int),
+  mesh_facetexcoord: wp.array[wp.vec3i],
+  mesh_texcoord: wp.array[wp.vec2],
+  mesh_texcoord_offsets: wp.array[int],
   hit_point: wp.vec3,
   bary_u: float,
   bary_v: float,
@@ -88,29 +88,29 @@ def sample_texture(
 @wp.func
 def cast_ray(
   # Model:
-  geom_type: wp.array(dtype=int),
-  geom_dataid: wp.array(dtype=int),
-  geom_size: wp.array2d(dtype=wp.vec3),
-  flex_vertadr: wp.array(dtype=int),
-  flex_edge: wp.array(dtype=wp.vec2i),
-  flex_radius: wp.array(dtype=float),
+  geom_type: wp.array[int],
+  geom_dataid: wp.array2d[int],
+  geom_size: wp.array2d[wp.vec3],
+  flex_vertadr: wp.array[int],
+  flex_edge: wp.array[wp.vec2i],
+  flex_radius: wp.array[float],
   # Data in:
-  geom_xpos_in: wp.array2d(dtype=wp.vec3),
-  geom_xmat_in: wp.array2d(dtype=wp.mat33),
-  flexvert_xpos_in: wp.array2d(dtype=wp.vec3),
+  geom_xpos_in: wp.array2d[wp.vec3],
+  geom_xmat_in: wp.array2d[wp.mat33],
+  flexvert_xpos_in: wp.array2d[wp.vec3],
   # In:
   bvh_id: wp.uint64,
   group_root: int,
   worldid: int,
   bvh_ngeom: int,
   flex_bvh_ngeom: int,
-  enabled_geom_ids: wp.array(dtype=int),
-  mesh_bvh_id: wp.array(dtype=wp.uint64),
-  hfield_bvh_id: wp.array(dtype=wp.uint64),
-  flex_geom_flexid: wp.array(dtype=int),
-  flex_geom_edgeid: wp.array(dtype=int),
-  flex_bvh_id: wp.array(dtype=wp.uint64),
-  flex_group_root: wp.array2d(dtype=int),
+  enabled_geom_ids: wp.array[int],
+  mesh_bvh_id: wp.array[wp.uint64],
+  hfield_bvh_id: wp.array[wp.uint64],
+  flex_geom_flexid: wp.array[int],
+  flex_geom_edgeid: wp.array[int],
+  flex_bvh_id: wp.array[wp.uint64],
+  flex_group_root: wp.array2d[int],
   ray_origin_world: wp.vec3,
   ray_dir_world: wp.vec3,
 ) -> Tuple[int, float, wp.vec3, float, float, int, int]:
@@ -159,7 +159,7 @@ def cast_ray(
     if gtype == GeomType.HFIELD:
       d, n, u, v, f, geom_hfield_id = ray_mesh_with_bvh(
         hfield_bvh_id,
-        geom_dataid[gi],
+        geom_dataid[worldid % geom_dataid.shape[0], gi],
         geom_xpos_in[worldid, gi],
         geom_xmat_in[worldid, gi],
         ray_origin_world,
@@ -208,7 +208,7 @@ def cast_ray(
     if gtype == GeomType.MESH:
       d, n, u, v, f, hit_mesh_id = ray_mesh_with_bvh(
         mesh_bvh_id,
-        geom_dataid[gi],
+        geom_dataid[worldid % geom_dataid.shape[0], gi],
         geom_xpos_in[worldid, gi],
         geom_xmat_in[worldid, gi],
         ray_origin_world,
@@ -256,29 +256,29 @@ def cast_ray(
 @wp.func
 def cast_ray_first_hit(
   # Model:
-  geom_type: wp.array(dtype=int),
-  geom_dataid: wp.array(dtype=int),
-  geom_size: wp.array2d(dtype=wp.vec3),
-  flex_vertadr: wp.array(dtype=int),
-  flex_edge: wp.array(dtype=wp.vec2i),
-  flex_radius: wp.array(dtype=float),
+  geom_type: wp.array[int],
+  geom_dataid: wp.array2d[int],
+  geom_size: wp.array2d[wp.vec3],
+  flex_vertadr: wp.array[int],
+  flex_edge: wp.array[wp.vec2i],
+  flex_radius: wp.array[float],
   # Data in:
-  geom_xpos_in: wp.array2d(dtype=wp.vec3),
-  geom_xmat_in: wp.array2d(dtype=wp.mat33),
-  flexvert_xpos_in: wp.array2d(dtype=wp.vec3),
+  geom_xpos_in: wp.array2d[wp.vec3],
+  geom_xmat_in: wp.array2d[wp.mat33],
+  flexvert_xpos_in: wp.array2d[wp.vec3],
   # In:
   bvh_id: wp.uint64,
   group_root: int,
   worldid: int,
   bvh_ngeom: int,
   bvh_nflexgeom: int,
-  enabled_geom_ids: wp.array(dtype=int),
-  mesh_bvh_id: wp.array(dtype=wp.uint64),
-  hfield_bvh_id: wp.array(dtype=wp.uint64),
-  flex_geom_flexid: wp.array(dtype=int),
-  flex_geom_edgeid: wp.array(dtype=int),
-  flex_bvh_id: wp.array(dtype=wp.uint64),
-  flex_group_root: wp.array2d(dtype=int),
+  enabled_geom_ids: wp.array[int],
+  mesh_bvh_id: wp.array[wp.uint64],
+  hfield_bvh_id: wp.array[wp.uint64],
+  flex_geom_flexid: wp.array[int],
+  flex_geom_edgeid: wp.array[int],
+  flex_bvh_id: wp.array[wp.uint64],
+  flex_group_root: wp.array2d[int],
   ray_origin_world: wp.vec3,
   ray_dir_world: wp.vec3,
   max_dist: float,
@@ -314,7 +314,7 @@ def cast_ray_first_hit(
     if gtype == GeomType.HFIELD:
       d, n, u, v, f, geom_hfield_id = ray_mesh_with_bvh(
         hfield_bvh_id,
-        geom_dataid[gi],
+        geom_dataid[worldid % geom_dataid.shape[0], gi],
         geom_xpos_in[worldid, gi],
         geom_xmat_in[worldid, gi],
         ray_origin_world,
@@ -363,7 +363,7 @@ def cast_ray_first_hit(
     if gtype == GeomType.MESH:
       hit = ray_mesh_with_bvh_anyhit(
         mesh_bvh_id,
-        geom_dataid[gi],
+        geom_dataid[worldid % geom_dataid.shape[0], gi],
         geom_xpos_in[worldid, gi],
         geom_xmat_in[worldid, gi],
         ray_origin_world,
@@ -409,30 +409,30 @@ def cast_ray_first_hit(
 @wp.func
 def compute_lighting(
   # Model:
-  geom_type: wp.array(dtype=int),
-  geom_dataid: wp.array(dtype=int),
-  geom_size: wp.array2d(dtype=wp.vec3),
-  flex_vertadr: wp.array(dtype=int),
-  flex_edge: wp.array(dtype=wp.vec2i),
-  flex_radius: wp.array(dtype=float),
+  geom_type: wp.array[int],
+  geom_dataid: wp.array2d[int],
+  geom_size: wp.array2d[wp.vec3],
+  flex_vertadr: wp.array[int],
+  flex_edge: wp.array[wp.vec2i],
+  flex_radius: wp.array[float],
   # Data in:
-  geom_xpos_in: wp.array2d(dtype=wp.vec3),
-  geom_xmat_in: wp.array2d(dtype=wp.mat33),
-  flexvert_xpos_in: wp.array2d(dtype=wp.vec3),
+  geom_xpos_in: wp.array2d[wp.vec3],
+  geom_xmat_in: wp.array2d[wp.mat33],
+  flexvert_xpos_in: wp.array2d[wp.vec3],
   # In:
   use_shadows: bool,
   bvh_id: wp.uint64,
   group_root: int,
   bvh_ngeom: int,
   bvh_nflexgeom: int,
-  enabled_geom_ids: wp.array(dtype=int),
+  enabled_geom_ids: wp.array[int],
   worldid: int,
-  mesh_bvh_id: wp.array(dtype=wp.uint64),
-  hfield_bvh_id: wp.array(dtype=wp.uint64),
-  flex_geom_flexid: wp.array(dtype=int),
-  flex_geom_edgeid: wp.array(dtype=int),
-  flex_bvh_id: wp.array(dtype=wp.uint64),
-  flex_group_root: wp.array2d(dtype=int),
+  mesh_bvh_id: wp.array[wp.uint64],
+  hfield_bvh_id: wp.array[wp.uint64],
+  flex_geom_flexid: wp.array[int],
+  flex_geom_edgeid: wp.array[int],
+  flex_bvh_id: wp.array[wp.uint64],
+  flex_group_root: wp.array2d[int],
   lightactive: bool,
   lighttype: int,
   lightcastshadow: bool,
@@ -530,65 +530,65 @@ def render(m: Model, d: Data, rc: RenderContext):
   @wp.kernel(module="unique", enable_backward=False)
   def _render_megakernel(
     # Model:
-    geom_type: wp.array(dtype=int),
-    geom_dataid: wp.array(dtype=int),
-    geom_matid: wp.array2d(dtype=int),
-    geom_size: wp.array2d(dtype=wp.vec3),
-    geom_rgba: wp.array2d(dtype=wp.vec4),
-    cam_projection: wp.array(dtype=int),
-    cam_fovy: wp.array2d(dtype=float),
-    cam_sensorsize: wp.array(dtype=wp.vec2),
-    cam_intrinsic: wp.array2d(dtype=wp.vec4),
-    light_type: wp.array2d(dtype=int),
-    light_castshadow: wp.array2d(dtype=bool),
-    light_active: wp.array2d(dtype=bool),
-    flex_vertadr: wp.array(dtype=int),
-    flex_edge: wp.array(dtype=wp.vec2i),
-    flex_radius: wp.array(dtype=float),
-    mesh_faceadr: wp.array(dtype=int),
-    mat_texid: wp.array3d(dtype=int),
-    mat_texrepeat: wp.array2d(dtype=wp.vec2),
-    mat_rgba: wp.array2d(dtype=wp.vec4),
+    geom_type: wp.array[int],
+    geom_dataid: wp.array2d[int],
+    geom_matid: wp.array2d[int],
+    geom_size: wp.array2d[wp.vec3],
+    geom_rgba: wp.array2d[wp.vec4],
+    cam_projection: wp.array[int],
+    cam_fovy: wp.array2d[float],
+    cam_sensorsize: wp.array[wp.vec2],
+    cam_intrinsic: wp.array2d[wp.vec4],
+    light_type: wp.array2d[int],
+    light_castshadow: wp.array2d[bool],
+    light_active: wp.array2d[bool],
+    flex_vertadr: wp.array[int],
+    flex_edge: wp.array[wp.vec2i],
+    flex_radius: wp.array[float],
+    mesh_faceadr: wp.array[int],
+    mat_texid: wp.array3d[int],
+    mat_texrepeat: wp.array2d[wp.vec2],
+    mat_rgba: wp.array2d[wp.vec4],
     # Data in:
-    geom_xpos_in: wp.array2d(dtype=wp.vec3),
-    geom_xmat_in: wp.array2d(dtype=wp.mat33),
-    cam_xpos_in: wp.array2d(dtype=wp.vec3),
-    cam_xmat_in: wp.array2d(dtype=wp.mat33),
-    light_xpos_in: wp.array2d(dtype=wp.vec3),
-    light_xdir_in: wp.array2d(dtype=wp.vec3),
-    flexvert_xpos_in: wp.array2d(dtype=wp.vec3),
+    geom_xpos_in: wp.array2d[wp.vec3],
+    geom_xmat_in: wp.array2d[wp.mat33],
+    cam_xpos_in: wp.array2d[wp.vec3],
+    cam_xmat_in: wp.array2d[wp.mat33],
+    light_xpos_in: wp.array2d[wp.vec3],
+    light_xdir_in: wp.array2d[wp.vec3],
+    flexvert_xpos_in: wp.array2d[wp.vec3],
     # In:
     nrender: int,
     use_shadows: bool,
     bvh_ngeom: int,
     bvh_nflexgeom: int,
-    cam_res: wp.array(dtype=wp.vec2i),
-    cam_id_map: wp.array(dtype=int),
-    ray: wp.array(dtype=wp.vec3),
-    rgb_adr: wp.array(dtype=int),
-    depth_adr: wp.array(dtype=int),
-    seg_adr: wp.array(dtype=int),
-    render_rgb: wp.array(dtype=bool),
-    render_depth: wp.array(dtype=bool),
-    render_seg: wp.array(dtype=bool),
+    cam_res: wp.array[wp.vec2i],
+    cam_id_map: wp.array[int],
+    ray: wp.array[wp.vec3],
+    rgb_adr: wp.array[int],
+    depth_adr: wp.array[int],
+    seg_adr: wp.array[int],
+    render_rgb: wp.array[bool],
+    render_depth: wp.array[bool],
+    render_seg: wp.array[bool],
     bvh_id: wp.uint64,
-    group_root: wp.array(dtype=int),
-    flex_bvh_id: wp.array(dtype=wp.uint64),
-    flex_group_root: wp.array2d(dtype=int),
-    enabled_geom_ids: wp.array(dtype=int),
-    mesh_bvh_id: wp.array(dtype=wp.uint64),
-    mesh_facetexcoord: wp.array(dtype=wp.vec3i),
-    mesh_texcoord: wp.array(dtype=wp.vec2),
-    mesh_texcoord_offsets: wp.array(dtype=int),
-    hfield_bvh_id: wp.array(dtype=wp.uint64),
-    flex_rgba: wp.array(dtype=wp.vec4),
-    flex_geom_flexid: wp.array(dtype=int),
-    flex_geom_edgeid: wp.array(dtype=int),
-    textures: wp.array(dtype=wp.Texture2D),
+    group_root: wp.array[int],
+    flex_bvh_id: wp.array[wp.uint64],
+    flex_group_root: wp.array2d[int],
+    enabled_geom_ids: wp.array[int],
+    mesh_bvh_id: wp.array[wp.uint64],
+    mesh_facetexcoord: wp.array[wp.vec3i],
+    mesh_texcoord: wp.array[wp.vec2],
+    mesh_texcoord_offsets: wp.array[int],
+    hfield_bvh_id: wp.array[wp.uint64],
+    flex_rgba: wp.array[wp.vec4],
+    flex_geom_flexid: wp.array[int],
+    flex_geom_edgeid: wp.array[int],
+    textures: wp.array[wp.Texture2D],
     # Out:
-    rgb_out: wp.array2d(dtype=wp.uint32),
-    depth_out: wp.array2d(dtype=float),
-    seg_out: wp.array2d(dtype=int),
+    rgb_out: wp.array2d[wp.uint32],
+    depth_out: wp.array2d[float],
+    seg_out: wp.array2d[int],
   ):
     worldid, rayid = wp.tid()
 

mjx/mujoco/mjx/third_party/mujoco_warp/_src/render_util.py

@@ -28,9 +28,9 @@ def _convert_texture_data(
   width: int,
   adr: int,
   nc: int,
-  tex_data_in: wp.array(dtype=wp.uint8),
+  tex_data_in: wp.array[wp.uint8],
   # Out:
-  tex_data_out: wp.array3d(dtype=float),
+  tex_data_out: wp.array3d[float],
 ):
   """Convert uint8 texture data to vec4 format for efficient sampling."""
   x, y = wp.tid()
@@ -134,11 +134,11 @@ def pack_rgba_to_uint32(r: float, g: float, b: float, a: float) -> wp.uint32:
 @wp.kernel
 def unpack_rgb_kernel(
   # In:
-  packed: wp.array2d(dtype=wp.uint32),
-  rgb_adr: wp.array(dtype=int),
+  packed: wp.array2d[wp.uint32],
+  rgb_adr: wp.array[int],
   camera_index: int,
   # Out:
-  rgb_out: wp.array3d(dtype=wp.vec3),
+  rgb_out: wp.array3d[wp.vec3],
 ):
   """Unpack ABGR uint32 packed pixel data into separate R, G, and B channels."""
   worldid, pixelid = wp.tid()
@@ -157,12 +157,12 @@ def unpack_rgb_kernel(
 @wp.kernel
 def extract_depth_kernel(
   # In:
-  depth_data: wp.array2d(dtype=float),
-  depth_adr: wp.array(dtype=int),
+  depth_data: wp.array2d[float],
+  depth_adr: wp.array[int],
   camera_index: int,
   depth_scale: float,
   # Out:
-  depth_out: wp.array3d(dtype=float),
+  depth_out: wp.array3d[float],
 ):
   """Extract the depth data from the render context buffers for a given camera index."""
   worldid, pixelid = wp.tid()
@@ -174,7 +174,7 @@ def extract_depth_kernel(
   depth_out[worldid, yid, xid] = wp.clamp(val / depth_scale, 0.0, 1.0)
 
 
-def get_rgb(rc: RenderContext, camera_index: int, rgb_out: wp.array3d(dtype=wp.vec3)):
+def get_rgb(rc: RenderContext, camera_index: int, rgb_out: wp.array3d[wp.vec3]):
   """Get the RGB data output from the render context buffers for a given camera index.
 
   Args:
@@ -190,7 +190,7 @@ def get_rgb(rc: RenderContext, camera_index: int, rgb_out: wp.array3d(dtype=wp.v
   )
 
 
-def get_depth(rc: RenderContext, camera_index: int, depth_scale: float, depth_out: wp.array3d(dtype=float)):
+def get_depth(rc: RenderContext, camera_index: int, depth_scale: float, depth_out: wp.array3d[float]):
   """Get the depth data output from the render context buffers for a given camera index.
 
   Args:
@@ -211,11 +211,11 @@ def get_depth(rc: RenderContext, camera_index: int, depth_scale: float, depth_ou
 @wp.kernel
 def _extract_seg_kernel(
   # In:
-  seg_data: wp.array2d(dtype=int),
-  seg_adr: wp.array(dtype=int),
+  seg_data: wp.array2d[int],
+  seg_adr: wp.array[int],
   camera_index: int,
   # Out:
-  seg_out: wp.array3d(dtype=int),
+  seg_out: wp.array3d[int],
 ):
   """Extract per-pixel geom IDs from the render context buffers for a given camera index."""
   worldid, pixelid = wp.tid()
@@ -226,7 +226,7 @@ def _extract_seg_kernel(
   seg_out[worldid, yid, xid] = seg_data[worldid, seg_adr_offset + pixelid]
 
 
-def get_segmentation(rc: RenderContext, camera_index: int, seg_out: wp.array3d(dtype=int)):
+def get_segmentation(rc: RenderContext, camera_index: int, seg_out: wp.array3d[int]):
   """Get the segmentation data from the render context buffers for a given camera index.
 
   Each pixel contains the MuJoCo geom ID of the geometry hit by the ray, -1 for

mjx/mujoco/mjx/third_party/mujoco_warp/_src/support.py

@@ -69,16 +69,16 @@ def mul_m_sparse(check_skip: bool):
   @wp.kernel(module="unique")
   def _mul_m_sparse(
     # Model:
-    qM_mulm_rowadr: wp.array(dtype=int),
-    qM_mulm_col: wp.array(dtype=int),
-    qM_mulm_madr: wp.array(dtype=int),
+    qM_mulm_rowadr: wp.array[int],
+    qM_mulm_col: wp.array[int],
+    qM_mulm_madr: wp.array[int],
     # Data in:
-    qM_in: wp.array3d(dtype=float),
+    qM_in: wp.array3d[float],
     # In:
-    vec: wp.array2d(dtype=float),
-    skip: wp.array(dtype=bool),
+    vec: wp.array2d[float],
+    skip: wp.array[bool],
     # Out:
-    res: wp.array2d(dtype=float),
+    res: wp.array2d[float],
   ):
     """Sparse matmul: one thread per DOF, gather-based (no atomics)."""
     worldid, dofid = wp.tid()
@@ -108,12 +108,12 @@ def mul_m_dense(nv: int, check_skip: bool):
   @wp.kernel(module="unique")
   def _mul_m_dense(
     # Data in:
-    qM_in: wp.array3d(dtype=float),
+    qM_in: wp.array3d[float],
     # In:
-    vec: wp.array2d(dtype=float),
-    skip: wp.array(dtype=bool),
+    vec: wp.array2d[float],
+    skip: wp.array[bool],
     # Out:
-    res: wp.array2d(dtype=float),
+    res: wp.array2d[float],
   ):
     worldid, i = wp.tid()
 
@@ -133,8 +133,8 @@ def mul_m_dense(nv: int, check_skip: bool):
 def mul_m(
   m: Model,
   d: Data,
-  res: wp.array2d(dtype=float),
-  vec: wp.array2d(dtype=float),
+  res: wp.array2d[float],
+  vec: wp.array2d[float],
   skip: Optional[wp.array] = None,
   M: Optional[wp.array] = None,
 ):
@@ -175,18 +175,18 @@ def mul_m(
 def _apply_ft(
   # Model:
   nbody: int,
-  body_parentid: wp.array(dtype=int),
-  body_rootid: wp.array(dtype=int),
-  dof_bodyid: wp.array(dtype=int),
+  body_parentid: wp.array[int],
+  body_rootid: wp.array[int],
+  dof_bodyid: wp.array[int],
   # Data in:
-  xipos_in: wp.array2d(dtype=wp.vec3),
-  subtree_com_in: wp.array2d(dtype=wp.vec3),
-  cdof_in: wp.array2d(dtype=wp.spatial_vector),
+  xipos_in: wp.array2d[wp.vec3],
+  subtree_com_in: wp.array2d[wp.vec3],
+  cdof_in: wp.array2d[wp.spatial_vector],
   # In:
-  ft_in: wp.array2d(dtype=wp.spatial_vector),
+  ft_in: wp.array2d[wp.spatial_vector],
   flg_add: bool,
   # Out:
-  qfrc_out: wp.array2d(dtype=float),
+  qfrc_out: wp.array2d[float],
 ):
   worldid, dofid = wp.tid()
   cdof = cdof_in[worldid, dofid]
@@ -216,7 +216,7 @@ def _apply_ft(
     qfrc_out[worldid, dofid] = accumul
 
 
-def apply_ft(m: Model, d: Data, ft: wp.array2d(dtype=wp.spatial_vector), qfrc: wp.array2d(dtype=float), flg_add: bool):
+def apply_ft(m: Model, d: Data, ft: wp.array2d[wp.spatial_vector], qfrc: wp.array2d[float], flg_add: bool):
   wp.launch(
     kernel=_apply_ft,
     dim=(d.nworld, m.nv),
@@ -226,7 +226,7 @@ def apply_ft(m: Model, d: Data, ft: wp.array2d(dtype=wp.spatial_vector), qfrc: w
 
 
 @event_scope
-def xfrc_accumulate(m: Model, d: Data, qfrc: wp.array2d(dtype=float)):
+def xfrc_accumulate(m: Model, d: Data, qfrc: wp.array2d[float]):
   """Map applied forces at each body via Jacobians to dof space and accumulate.
 
   Args:
@@ -238,9 +238,7 @@ def xfrc_accumulate(m: Model, d: Data, qfrc: wp.array2d(dtype=float)):
 
 
 @wp.func
-def _decode_pyramid(
-  njmax_in: int, pyramid: wp.array(dtype=float), efc_address: int, mu: vec5, condim: int
-) -> wp.spatial_vector:
+def _decode_pyramid(njmax_in: int, pyramid: wp.array[float], efc_address: int, mu: vec5, condim: int) -> wp.spatial_vector:
   """Converts pyramid representation to contact force."""
   force = wp.spatial_vector()
 
@@ -270,13 +268,13 @@ def contact_force_fn(
   # Model:
   opt_cone: int,
   # Data in:
-  contact_frame_in: wp.array(dtype=wp.mat33),
-  contact_friction_in: wp.array(dtype=vec5),
-  contact_dim_in: wp.array(dtype=int),
-  contact_efc_address_in: wp.array2d(dtype=int),
-  efc_force_in: wp.array2d(dtype=float),
+  contact_frame_in: wp.array[wp.mat33],
+  contact_friction_in: wp.array[vec5],
+  contact_dim_in: wp.array[int],
+  contact_efc_address_in: wp.array2d[int],
+  efc_force_in: wp.array2d[float],
   njmax_in: int,
-  nacon_in: wp.array(dtype=int),
+  nacon_in: wp.array[int],
   # In:
   worldid: int,
   contact_id: int,
@@ -315,19 +313,19 @@ def contact_force_kernel(
   # Model:
   opt_cone: int,
   # Data in:
-  contact_frame_in: wp.array(dtype=wp.mat33),
-  contact_friction_in: wp.array(dtype=vec5),
-  contact_dim_in: wp.array(dtype=int),
-  contact_efc_address_in: wp.array2d(dtype=int),
-  contact_worldid_in: wp.array(dtype=int),
-  efc_force_in: wp.array2d(dtype=float),
+  contact_frame_in: wp.array[wp.mat33],
+  contact_friction_in: wp.array[vec5],
+  contact_dim_in: wp.array[int],
+  contact_efc_address_in: wp.array2d[int],
+  contact_worldid_in: wp.array[int],
+  efc_force_in: wp.array2d[float],
   njmax_in: int,
-  nacon_in: wp.array(dtype=int),
+  nacon_in: wp.array[int],
   # In:
-  contact_ids: wp.array(dtype=int),
+  contact_ids: wp.array[int],
   to_world_frame: bool,
   # Out:
-  out: wp.array(dtype=wp.spatial_vector),
+  out: wp.array[wp.spatial_vector],
 ):
   tid = wp.tid()
 
@@ -353,9 +351,7 @@ def contact_force_kernel(
   )
 
 
-def contact_force(
-  m: Model, d: Data, contact_ids: wp.array(dtype=int), to_world_frame: bool, force: wp.array(dtype=wp.spatial_vector)
-):
+def contact_force(m: Model, d: Data, contact_ids: wp.array[int], to_world_frame: bool, force: wp.array[wp.spatial_vector]):
   """Compute forces for contacts in Data.
 
   Args:
@@ -400,12 +396,12 @@ def transform_force(frc: wp.spatial_vector, offset: wp.vec3) -> wp.spatial_vecto
 @wp.func
 def jac_dof(
   # Model:
-  body_parentid: wp.array(dtype=int),
-  body_rootid: wp.array(dtype=int),
-  dof_bodyid: wp.array(dtype=int),
+  body_parentid: wp.array[int],
+  body_rootid: wp.array[int],
+  dof_bodyid: wp.array[int],
   # Data in:
-  subtree_com_in: wp.array2d(dtype=wp.vec3),
-  cdof_in: wp.array2d(dtype=wp.spatial_vector),
+  subtree_com_in: wp.array2d[wp.vec3],
+  cdof_in: wp.array2d[wp.spatial_vector],
   # In:
   point: wp.vec3,
   bodyid: int,
@@ -441,18 +437,18 @@ def _make_jac_kernel(has_jacp: bool, has_jacr: bool):
   @wp.kernel(module="unique", enable_backward=False)
   def _jac(
     # Model:
-    body_parentid: wp.array(dtype=int),
-    body_rootid: wp.array(dtype=int),
-    dof_bodyid: wp.array(dtype=int),
+    body_parentid: wp.array[int],
+    body_rootid: wp.array[int],
+    dof_bodyid: wp.array[int],
     # Data in:
-    subtree_com_in: wp.array2d(dtype=wp.vec3),
-    cdof_in: wp.array2d(dtype=wp.spatial_vector),
+    subtree_com_in: wp.array2d[wp.vec3],
+    cdof_in: wp.array2d[wp.spatial_vector],
     # In:
-    point_in: wp.array(dtype=wp.vec3),
-    bodyid_in: wp.array(dtype=int),
+    point_in: wp.array[wp.vec3],
+    bodyid_in: wp.array[int],
     # Out:
-    jacp_out: wp.array3d(dtype=float),
-    jacr_out: wp.array3d(dtype=float),
+    jacp_out: wp.array3d[float],
+    jacr_out: wp.array3d[float],
   ):
     worldid, dofid = wp.tid()
 
@@ -477,10 +473,10 @@ def _make_jac_kernel(has_jacp: bool, has_jacr: bool):
 def jac(
   m: Model,
   d: Data,
-  jacp: wp.array | None,  # wp.array3d(dtype=float)
-  jacr: wp.array | None,  # wp.array3d(dtype=float)
-  point: wp.array(dtype=wp.vec3),
-  body: wp.array(dtype=int),
+  jacp: wp.array | None,  # wp.array3d[float]
+  jacr: wp.array | None,  # wp.array3d[float]
+  point: wp.array[wp.vec3],
+  body: wp.array[int],
 ):
   """Compute translational and rotational Jacobian for point on body.
 
@@ -508,17 +504,17 @@ def jac(
 @wp.func
 def jac_dot_dof(
   # Model:
-  body_parentid: wp.array(dtype=int),
-  body_rootid: wp.array(dtype=int),
-  jnt_type: wp.array(dtype=int),
-  jnt_dofadr: wp.array(dtype=int),
-  dof_bodyid: wp.array(dtype=int),
-  dof_jntid: wp.array(dtype=int),
+  body_parentid: wp.array[int],
+  body_rootid: wp.array[int],
+  jnt_type: wp.array[int],
+  jnt_dofadr: wp.array[int],
+  dof_bodyid: wp.array[int],
+  dof_jntid: wp.array[int],
   # Data in:
-  subtree_com_in: wp.array2d(dtype=wp.vec3),
-  cdof_in: wp.array2d(dtype=wp.spatial_vector),
-  cvel_in: wp.array2d(dtype=wp.spatial_vector),
-  cdof_dot_in: wp.array2d(dtype=wp.spatial_vector),
+  subtree_com_in: wp.array2d[wp.vec3],
+  cdof_in: wp.array2d[wp.spatial_vector],
+  cvel_in: wp.array2d[wp.spatial_vector],
+  cdof_dot_in: wp.array2d[wp.spatial_vector],
   # In:
   point: wp.vec3,
   bodyid: int,
@@ -573,7 +569,7 @@ def jac_dot_dof(
   return jacp, jacr
 
 
-def get_state(m: Model, d: Data, state: wp.array2d(dtype=float), sig: int, active: Optional[wp.array] = None):
+def get_state(m: Model, d: Data, state: wp.array2d[float], sig: int, active: Optional[wp.array] = None):
   """Copy concatenated state components specified by sig from Data into state.
 
   The bits of the integer sig correspond to element fields of State.
@@ -599,22 +595,22 @@ def get_state(m: Model, d: Data, state: wp.array2d(dtype=float), sig: int, activ
     neq: int,
     nmocap: int,
     # Data in:
-    time_in: wp.array(dtype=float),
-    qpos_in: wp.array2d(dtype=float),
-    qvel_in: wp.array2d(dtype=float),
-    act_in: wp.array2d(dtype=float),
-    qacc_warmstart_in: wp.array2d(dtype=float),
-    ctrl_in: wp.array2d(dtype=float),
-    qfrc_applied_in: wp.array2d(dtype=float),
-    xfrc_applied_in: wp.array2d(dtype=wp.spatial_vector),
-    eq_active_in: wp.array2d(dtype=bool),
-    mocap_pos_in: wp.array2d(dtype=wp.vec3),
-    mocap_quat_in: wp.array2d(dtype=wp.quat),
+    time_in: wp.array[float],
+    qpos_in: wp.array2d[float],
+    qvel_in: wp.array2d[float],
+    act_in: wp.array2d[float],
+    qacc_warmstart_in: wp.array2d[float],
+    ctrl_in: wp.array2d[float],
+    qfrc_applied_in: wp.array2d[float],
+    xfrc_applied_in: wp.array2d[wp.spatial_vector],
+    eq_active_in: wp.array2d[bool],
+    mocap_pos_in: wp.array2d[wp.vec3],
+    mocap_quat_in: wp.array2d[wp.quat],
     # In:
     sig_in: int,
-    active_in: wp.array(dtype=bool),
+    active_in: wp.array[bool],
     # Out:
-    state_out: wp.array2d(dtype=float),
+    state_out: wp.array2d[float],
   ):
     worldid = wp.tid()
 
@@ -712,7 +708,7 @@ def get_state(m: Model, d: Data, state: wp.array2d(dtype=float), sig: int, activ
   )
 
 
-def set_state(m: Model, d: Data, state: wp.array2d(dtype=float), sig: int, active: Optional[wp.array] = None):
+def set_state(m: Model, d: Data, state: wp.array2d[float], sig: int, active: Optional[wp.array] = None):
   """Copy concatenated state components specified by sig from state into Data.
 
   The bits of the integer sig correspond to element fields of State.
@@ -739,20 +735,20 @@ def set_state(m: Model, d: Data, state: wp.array2d(dtype=float), sig: int, activ
     nmocap: int,
     # In:
     sig_in: int,
-    active_in: wp.array(dtype=bool),
-    state_in: wp.array2d(dtype=float),
+    active_in: wp.array[bool],
+    state_in: wp.array2d[float],
     # Data out:
-    time_out: wp.array(dtype=float),
-    qpos_out: wp.array2d(dtype=float),
-    qvel_out: wp.array2d(dtype=float),
-    act_out: wp.array2d(dtype=float),
-    qacc_warmstart_out: wp.array2d(dtype=float),
-    ctrl_out: wp.array2d(dtype=float),
-    qfrc_applied_out: wp.array2d(dtype=float),
-    xfrc_applied_out: wp.array2d(dtype=wp.spatial_vector),
-    eq_active_out: wp.array2d(dtype=bool),
-    mocap_pos_out: wp.array2d(dtype=wp.vec3),
-    mocap_quat_out: wp.array2d(dtype=wp.quat),
+    time_out: wp.array[float],
+    qpos_out: wp.array2d[float],
+    qvel_out: wp.array2d[float],
+    act_out: wp.array2d[float],
+    qacc_warmstart_out: wp.array2d[float],
+    ctrl_out: wp.array2d[float],
+    qfrc_applied_out: wp.array2d[float],
+    xfrc_applied_out: wp.array2d[wp.spatial_vector],
+    eq_active_out: wp.array2d[bool],
+    mocap_pos_out: wp.array2d[wp.vec3],
+    mocap_quat_out: wp.array2d[wp.quat],
   ):
     worldid = wp.tid()
 

mjx/mujoco/mjx/third_party/mujoco_warp/_src/types.py

@@ -17,6 +17,7 @@ import enum
 from typing import Callable
 
 import mujoco
+import numpy as np
 import warp as wp
 
 MJ_MINVAL = mujoco.mjMINVAL
@@ -793,9 +794,18 @@ class TileSet:
     size: size of all the tiles in this set
   """
 
-  adr: wp.array(dtype=int)
+  adr: wp.array[int]
   size: int
 
+  def __eq__(self, other) -> bool:
+    if self.__class__ is not other.__class__:
+      return NotImplemented
+    return self.size == other.size and np.array_equal(np.asarray(self.adr.numpy()), np.asarray(other.adr.numpy()))
+
+  def __hash__(self) -> int:
+    adr = np.asarray(self.adr.numpy())
+    return hash((self.size, adr.dtype.str, adr.shape, adr.tobytes()))
+
 
 @dataclasses.dataclass
 class Callback:
@@ -933,7 +943,7 @@ class Model:
     geom_conaffinity: geom contact affinity                  (ngeom,)
     geom_condim: contact dimensionality (1, 3, 4, 6)         (ngeom,)
     geom_bodyid: id of geom's body                           (ngeom,)
-    geom_dataid: id of geom's mesh/hfield; -1: none          (ngeom,)
+    geom_dataid: id of geom's mesh/hfield; -1: none          (*, ngeom)
     geom_matid: material id for rendering                    (*, ngeom,)
     geom_group: geom group inclusion/exclusion mask          (ngeom,)
     geom_priority: geom contact priority                     (ngeom,)
@@ -1322,7 +1332,7 @@ class Model:
   geom_conaffinity: array("ngeom", int)
   geom_condim: array("ngeom", int)
   geom_bodyid: array("ngeom", int)
-  geom_dataid: array("ngeom", int)
+  geom_dataid: array("*", "ngeom", int)
   geom_matid: array("*", "ngeom", int)
   geom_group: array("ngeom", int)
   geom_priority: array("ngeom", int)
@@ -1527,70 +1537,70 @@ class Model:
   has_fluid: bool
   has_sdf_geom: bool
   block_dim: BlockDim
-  body_tree: tuple[wp.array(dtype=int), ...]
-  body_branches: wp.array(dtype=int)
-  body_branch_start: wp.array(dtype=int)
+  body_tree: tuple[wp.array[int], ...]
+  body_branches: wp.array[int]
+  body_branch_start: wp.array[int]
   mocap_bodyid: array("nmocap", int)
   body_fluid_ellipsoid: array("nbody", bool)
-  jnt_limited_slide_hinge_adr: wp.array(dtype=int)
-  jnt_limited_ball_adr: wp.array(dtype=int)
-  dof_tri_row: wp.array(dtype=int)
-  dof_tri_col: wp.array(dtype=int)
-  nxn_geom_pair: wp.array(dtype=wp.vec2i)
-  nxn_geom_pair_filtered: wp.array(dtype=wp.vec2i)
-  nxn_pairid: wp.array(dtype=wp.vec2i)
-  nxn_pairid_filtered: wp.array(dtype=wp.vec2i)
+  jnt_limited_slide_hinge_adr: wp.array[int]
+  jnt_limited_ball_adr: wp.array[int]
+  dof_tri_row: wp.array[int]
+  dof_tri_col: wp.array[int]
+  nxn_geom_pair: wp.array[wp.vec2i]
+  nxn_geom_pair_filtered: wp.array[wp.vec2i]
+  nxn_pairid: wp.array[wp.vec2i]
+  nxn_pairid_filtered: wp.array[wp.vec2i]
   geom_pair_type_count: tuple[int, ...]
   geom_plugin_index: array("ngeom", int)
-  eq_connect_adr: wp.array(dtype=int)
-  eq_wld_adr: wp.array(dtype=int)
-  eq_jnt_adr: wp.array(dtype=int)
-  eq_ten_adr: wp.array(dtype=int)
-  eq_flex_adr: wp.array(dtype=int)
-  tendon_jnt_adr: wp.array(dtype=int)
-  tendon_site_pair_adr: wp.array(dtype=int)
-  tendon_geom_adr: wp.array(dtype=int)
-  tendon_limited_adr: wp.array(dtype=int)
+  eq_connect_adr: wp.array[int]
+  eq_wld_adr: wp.array[int]
+  eq_jnt_adr: wp.array[int]
+  eq_ten_adr: wp.array[int]
+  eq_flex_adr: wp.array[int]
+  tendon_jnt_adr: wp.array[int]
+  tendon_site_pair_adr: wp.array[int]
+  tendon_geom_adr: wp.array[int]
+  tendon_limited_adr: wp.array[int]
   max_ten_J_rownnz: int
-  ten_wrapadr_site: wp.array(dtype=int)
-  ten_wrapnum_site: wp.array(dtype=int)
-  wrap_jnt_adr: wp.array(dtype=int)
-  wrap_site_adr: wp.array(dtype=int)
-  wrap_site_pair_adr: wp.array(dtype=int)
-  wrap_geom_adr: wp.array(dtype=int)
+  ten_wrapadr_site: wp.array[int]
+  ten_wrapnum_site: wp.array[int]
+  wrap_jnt_adr: wp.array[int]
+  wrap_site_adr: wp.array[int]
+  wrap_site_pair_adr: wp.array[int]
+  wrap_geom_adr: wp.array[int]
   wrap_pulley_scale: array("nwrap", float)
-  actuator_trntype_body_adr: wp.array(dtype=int)
-  sensor_pos_adr: wp.array(dtype=int)
-  sensor_limitpos_adr: wp.array(dtype=int)
-  sensor_vel_adr: wp.array(dtype=int)
-  sensor_limitvel_adr: wp.array(dtype=int)
-  sensor_acc_adr: wp.array(dtype=int)
-  sensor_rangefinder_adr: wp.array(dtype=int)
-  rangefinder_sensor_adr: wp.array(dtype=int)
-  sensor_collision_start_adr: wp.array(dtype=int)
+  actuator_trntype_body_adr: wp.array[int]
+  sensor_pos_adr: wp.array[int]
+  sensor_limitpos_adr: wp.array[int]
+  sensor_vel_adr: wp.array[int]
+  sensor_limitvel_adr: wp.array[int]
+  sensor_acc_adr: wp.array[int]
+  sensor_rangefinder_adr: wp.array[int]
+  rangefinder_sensor_adr: wp.array[int]
+  sensor_collision_start_adr: wp.array[int]
   collision_sensor_adr: array("nsensor", int)
-  sensor_touch_adr: wp.array(dtype=int)
-  sensor_limitfrc_adr: wp.array(dtype=int)
+  sensor_touch_adr: wp.array[int]
+  sensor_limitfrc_adr: wp.array[int]
   sensor_e_potential: bool
   sensor_e_kinetic: bool
-  sensor_tendonactfrc_adr: wp.array(dtype=int)
+  sensor_tendonactfrc_adr: wp.array[int]
   sensor_subtree_vel: bool
   sensor_contact_adr: array("nsensorcontact", int)
   sensor_adr_to_contact_adr: array("nsensor", int)
   sensor_rne_postconstraint: bool
   sensor_rangefinder_bodyid: array("nrangefinder", int)
   taxel_vertadr: array("nsensortaxel", int)
-  taxel_sensorid: wp.array(dtype=int)
+  taxel_sensorid: wp.array[int]
   qM_tiles: tuple[TileSet, ...]
-  qLD_updates: tuple[wp.array(dtype=wp.vec3i), ...]
-  qLD_all_updates: wp.array(dtype=wp.vec3i)
-  qLD_level_offsets: wp.array(dtype=int)
-  qM_fullm_i: wp.array(dtype=int)
-  qM_fullm_j: wp.array(dtype=int)
+  qLD_updates: tuple[wp.array[wp.vec3i], ...]
+  qLD_all_updates: wp.array[wp.vec3i]
+  qLD_level_offsets: wp.array[int]
+  qM_fullm_i: wp.array[int]
+  qM_fullm_j: wp.array[int]
   # Gather-based sparse mul_m indices (thread per DOF, no atomics)
-  qM_mulm_rowadr: wp.array(dtype=int)  # start address for each row [nv+1]
-  qM_mulm_col: wp.array(dtype=int)  # column index to gather from
-  qM_mulm_madr: wp.array(dtype=int)  # matrix address to read
+  qM_mulm_rowadr: wp.array[int]  # start address for each row [nv+1]
+  qM_mulm_col: wp.array[int]  # column index to gather from
+  qM_mulm_madr: wp.array[int]  # matrix address to read
 
 
 class ContactType(enum.IntFlag):
@@ -1674,10 +1684,10 @@ class Constraint:
 
   type: array("nworld", "njmax", int)
   id: array("nworld", "njmax", int)
-  J_rownnz: wp.array2d(dtype=int)
-  J_rowadr: wp.array2d(dtype=int)
-  J_colind: wp.array3d(dtype=int)
-  J: wp.array3d(dtype=float)
+  J_rownnz: wp.array2d[int]
+  J_rowadr: wp.array2d[int]
+  J_colind: wp.array3d[int]
+  J: wp.array3d[float]
   pos: array("nworld", "njmax", float)
   margin: array("nworld", "njmax", float)
   D: array("nworld", "njmax_pad", float)
@@ -1846,8 +1856,8 @@ class Data:
   moment_colind: array("nworld", "nJmom", int)
   actuator_moment: array("nworld", "nJmom", float)
   crb: array("nworld", "nbody", vec10)
-  qM: wp.array3d(dtype=float)
-  qLD: wp.array3d(dtype=float)
+  qM: wp.array3d[float]
+  qLD: wp.array3d[float]
   qLDiagInv: array("nworld", "nv", float)
   flexedge_velocity: array("nworld", "nflexedge", float)
   ten_velocity: array("nworld", "ntendon", float)

mjx/mujoco/mjx/third_party/mujoco_warp/pyproject.toml

@@ -57,6 +57,7 @@ dev = [
     "lsprotocol>=2023.0.1,<2024.0.0",
     "mujoco>=3.6.0.dev0",
     "warp-lang>=1.11.0.dev0",
+    "mjviser>=0.0.10",
 ]
 # TODO(team): cpu and cuda JAX optional dependencies are temporary, remove after we land MJX:Warp
 cpu = [

mjx/mujoco/mjx/third_party/mujoco_warp/viewer.py

@@ -1,4 +1,4 @@
-# Copyright 2025 The Newton Developers
+# Copyright 2026 The Newton Developers
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.
@@ -62,6 +62,7 @@ _OVERRIDE = flags.DEFINE_multi_string("override", [], "Model overrides (notation
 _KEYFRAME = flags.DEFINE_integer("keyframe", 0, "keyframe to initialize simulation.")
 _DEVICE = flags.DEFINE_string("device", None, "override the default Warp device")
 _REPLAY = flags.DEFINE_string("replay", None, "keyframe sequence to replay, keyframe name must prefix match")
+_VIEWER = flags.DEFINE_enum("viewer", "mujoco", ["mujoco", "viser"], "Viewer backend (mujoco native or mjviser web)")
 
 _VIEWER_GLOBAL_STATE = {"running": True, "step_once": False}
 
@@ -105,6 +106,70 @@ def _compile_step(m, d):
   return capture.graph
 
 
+def _make_warp_step_fn(mjm, m, d, graph, ctrls=None):
+  ctrlid = 0
+  opt = copy.copy(mjm.opt)
+
+  def step_fn(mjm, mjd):
+    nonlocal ctrlid, opt, m, graph
+    if ctrls is not None and ctrlid < len(ctrls):
+      mjd.ctrl[:] = ctrls[ctrlid]
+      ctrlid += 1
+    if mjm.opt != opt:
+      opt = copy.copy(mjm.opt)
+      m = mjw.put_model(mjm)
+      graph = _compile_step(m, d) if wp.get_device().is_cuda else None
+    wp.copy(d.ctrl, wp.array([mjd.ctrl.astype(np.float32)]))
+    wp.copy(d.act, wp.array([mjd.act.astype(np.float32)]))
+    wp.copy(d.xfrc_applied, wp.array([mjd.xfrc_applied.astype(np.float32)]))
+    wp.copy(d.qpos, wp.array([mjd.qpos.astype(np.float32)]))
+    wp.copy(d.qvel, wp.array([mjd.qvel.astype(np.float32)]))
+    wp.copy(d.time, wp.array([mjd.time], dtype=wp.float32))
+    if graph is None:
+      mjw.step(m, d)
+    else:
+      wp.capture_launch(graph)
+      wp.synchronize()
+    mjw.get_data_into(mjd, mjm, d)
+
+  return step_fn
+
+
+def _make_c_step_fn(ctrls=None):
+  if ctrls is None:
+    return mujoco.mj_step
+
+  ctrlid = 0
+
+  def step_fn(mjm, mjd):
+    nonlocal ctrlid
+    if ctrlid < len(ctrls):
+      mjd.ctrl[:] = ctrls[ctrlid]
+      ctrlid += 1
+    mujoco.mj_step(mjm, mjd)
+
+  return step_fn
+
+
+def _run_viser_viewer(mjm, mjd, step_fn):
+  from mjviser import Viewer as MjViserViewer
+
+  MjViserViewer(mjm, mjd, step_fn=step_fn).run()
+
+
+def _run_passive_viewer(mjm, mjd, step_fn):
+  with mujoco.viewer.launch_passive(mjm, mjd, key_callback=key_callback) as viewer:
+    while True:
+      start = time.time()
+      if _VIEWER_GLOBAL_STATE["running"] or _VIEWER_GLOBAL_STATE["step_once"]:
+        _VIEWER_GLOBAL_STATE["step_once"] = False
+        step_fn(mjm, mjd)
+      viewer.sync()
+      elapsed = time.time() - start
+      if elapsed < mjm.opt.timestep:
+        time.sleep(mjm.opt.timestep - elapsed)
+
+
 def _main(argv: Sequence[str]) -> None:
   """Runs viewer app."""
   if len(argv) < 2:
@@ -115,7 +180,6 @@ def _main(argv: Sequence[str]) -> None:
   mjm = _load_model(epath.Path(argv[1]))
   mjd = mujoco.MjData(mjm)
   ctrls = None
-  ctrlid = 0
   if _REPLAY.value:
     keys = find_keys(mjm, _REPLAY.value)
     if not keys:
@@ -169,45 +233,15 @@ def _main(argv: Sequence[str]) -> None:
     print(f"Data\n  nworld: {d.nworld} nconmax: {int(d.naconmax / d.nworld)} njmax: {d.njmax}\n")
     print(f"MuJoCo Warp simulating with dt = {m.opt.timestep.numpy()[0]:.3f}...")
 
-  with mujoco.viewer.launch_passive(mjm, mjd, key_callback=key_callback) as viewer:
-    opt = copy.copy(mjm.opt)
+  if _ENGINE.value == EngineOptions.WARP:
+    step_fn = _make_warp_step_fn(mjm, m, d, graph, ctrls)
+  else:
+    step_fn = _make_c_step_fn(ctrls)
 
-    while True:
-      start = time.time()
-
-      if ctrls is not None and ctrlid < len(ctrls):
-        mjd.ctrl[:] = ctrls[ctrlid]
-        ctrlid += 1
-
-      if _ENGINE.value == EngineOptions.C:
-        mujoco.mj_step(mjm, mjd)
-      else:  # mjwarp
-        wp.copy(d.ctrl, wp.array([mjd.ctrl.astype(np.float32)]))
-        wp.copy(d.act, wp.array([mjd.act.astype(np.float32)]))
-        wp.copy(d.xfrc_applied, wp.array([mjd.xfrc_applied.astype(np.float32)]))
-        wp.copy(d.qpos, wp.array([mjd.qpos.astype(np.float32)]))
-        wp.copy(d.qvel, wp.array([mjd.qvel.astype(np.float32)]))
-        wp.copy(d.time, wp.array([mjd.time], dtype=wp.float32))
-        # if the user changed an option in the MuJoCo Simulate UI, go ahead and recompile the step
-        # TODO: update memory tied to option max iterations
-        if mjm.opt != opt:
-          opt = copy.copy(mjm.opt)
-          m = mjw.put_model(mjm)
-          graph = _compile_step(m, d) if wp.get_device().is_cuda else None
-        if _VIEWER_GLOBAL_STATE["running"] or _VIEWER_GLOBAL_STATE["step_once"]:
-          _VIEWER_GLOBAL_STATE["step_once"] = False
-          if graph is None:
-            mjw.step(m, d)
-          else:
-            wp.capture_launch(graph)
-            wp.synchronize()
-        mjw.get_data_into(mjd, mjm, d)
-
-      viewer.sync()
-
-      elapsed = time.time() - start
-      if elapsed < mjm.opt.timestep:
-        time.sleep(mjm.opt.timestep - elapsed)
+  if _VIEWER.value == "viser":
+    _run_viser_viewer(mjm, mjd, step_fn)
+  else:
+    _run_passive_viewer(mjm, mjd, step_fn)
 
 
 def main():

mjx/mujoco/mjx/warp/bvh.py

@@ -48,34 +48,35 @@ _cb = mjwp_types.Callback(
     **{f.name: None for f in dataclasses.fields(mjwp_types.Callback) if f.init}
 )
 
+
 @ffi.format_args_for_warp
 def _refit_bvh_shim(
     # Model
     nworld: int,
-    flex_dim: wp.array(dtype=int),
-    flex_edge: wp.array(dtype=wp.vec2i),
-    flex_elem: wp.array(dtype=int),
-    flex_elemadr: wp.array(dtype=int),
-    flex_elemdataadr: wp.array(dtype=int),
-    flex_elemnum: wp.array(dtype=int),
-    flex_radius: wp.array(dtype=float),
-    flex_shell: wp.array(dtype=int),
-    flex_shelldataadr: wp.array(dtype=int),
-    flex_vertadr: wp.array(dtype=int),
-    flex_vertnum: wp.array(dtype=int),
-    geom_dataid: wp.array(dtype=int),
-    geom_size: wp.array2d(dtype=wp.vec3),
-    geom_type: wp.array(dtype=int),
+    flex_dim: wp.array[int],
+    flex_edge: wp.array[wp.vec2i],
+    flex_elem: wp.array[int],
+    flex_elemadr: wp.array[int],
+    flex_elemdataadr: wp.array[int],
+    flex_elemnum: wp.array[int],
+    flex_radius: wp.array[float],
+    flex_shell: wp.array[int],
+    flex_shelldataadr: wp.array[int],
+    flex_vertadr: wp.array[int],
+    flex_vertnum: wp.array[int],
+    geom_dataid: wp.array2d[int],
+    geom_size: wp.array2d[wp.vec3],
+    geom_type: wp.array[int],
     nflex: int,
     nflexelem: int,
     # Data
-    flexvert_xpos: wp.array2d(dtype=wp.vec3),
-    geom_xmat: wp.array2d(dtype=wp.mat33),
-    geom_xpos: wp.array2d(dtype=wp.vec3),
+    flexvert_xpos: wp.array2d[wp.vec3],
+    geom_xmat: wp.array2d[wp.mat33],
+    geom_xpos: wp.array2d[wp.vec3],
     # Registry
     rc_id: int,
     # Dummy output
-    dummy: wp.array(dtype=int),
+    dummy: wp.array[int],
 ):
   _m.stat = _s
   _m.opt = _o
@@ -135,7 +136,7 @@ def _refit_bvh_jax_impl(
       m._impl.flex_shelldataadr,
       m.flex_vertadr,
       m.flex_vertnum,
-      m.geom_dataid,
+      jax.numpy.expand_dims(m.geom_dataid, 0),
       m.geom_size,
       m.geom_type,
       m.nflex,

mjx/mujoco/mjx/warp/collision_driver.py

@@ -46,67 +46,68 @@ _cb = mjwp_types.Callback(
     **{f.name: None for f in dataclasses.fields(mjwp_types.Callback) if f.init}
 )
 
+
 @ffi.format_args_for_warp
 def _collision_shim(
     # Model
     nworld: int,
     block_dim: mjwp_types.BlockDim,
-    flex_conaffinity: wp.array(dtype=int),
-    flex_condim: wp.array(dtype=int),
-    flex_contype: wp.array(dtype=int),
-    flex_dim: wp.array(dtype=int),
-    flex_elem: wp.array(dtype=int),
-    flex_elemadr: wp.array(dtype=int),
-    flex_elemdataadr: wp.array(dtype=int),
-    flex_elemnum: wp.array(dtype=int),
-    flex_friction: wp.array(dtype=wp.vec3),
-    flex_margin: wp.array(dtype=float),
-    flex_radius: wp.array(dtype=float),
-    flex_shell: wp.array(dtype=int),
-    flex_shelldataadr: wp.array(dtype=int),
-    flex_shellnum: wp.array(dtype=int),
-    flex_vertadr: wp.array(dtype=int),
-    flex_vertflexid: wp.array(dtype=int),
-    geom_aabb: wp.array3d(dtype=wp.vec3),
-    geom_conaffinity: wp.array(dtype=int),
-    geom_condim: wp.array(dtype=int),
-    geom_contype: wp.array(dtype=int),
-    geom_dataid: wp.array(dtype=int),
-    geom_friction: wp.array2d(dtype=wp.vec3),
-    geom_gap: wp.array2d(dtype=float),
-    geom_margin: wp.array2d(dtype=float),
+    flex_conaffinity: wp.array[int],
+    flex_condim: wp.array[int],
+    flex_contype: wp.array[int],
+    flex_dim: wp.array[int],
+    flex_elem: wp.array[int],
+    flex_elemadr: wp.array[int],
+    flex_elemdataadr: wp.array[int],
+    flex_elemnum: wp.array[int],
+    flex_friction: wp.array[wp.vec3],
+    flex_margin: wp.array[float],
+    flex_radius: wp.array[float],
+    flex_shell: wp.array[int],
+    flex_shelldataadr: wp.array[int],
+    flex_shellnum: wp.array[int],
+    flex_vertadr: wp.array[int],
+    flex_vertflexid: wp.array[int],
+    geom_aabb: wp.array3d[wp.vec3],
+    geom_conaffinity: wp.array[int],
+    geom_condim: wp.array[int],
+    geom_contype: wp.array[int],
+    geom_dataid: wp.array2d[int],
+    geom_friction: wp.array2d[wp.vec3],
+    geom_gap: wp.array2d[float],
+    geom_margin: wp.array2d[float],
     geom_pair_type_count: tuple[int, ...],
-    geom_plugin_index: wp.array(dtype=int),
-    geom_priority: wp.array(dtype=int),
-    geom_rbound: wp.array2d(dtype=float),
-    geom_size: wp.array2d(dtype=wp.vec3),
-    geom_solimp: wp.array2d(dtype=mjwp_types.vec5),
-    geom_solmix: wp.array2d(dtype=float),
-    geom_solref: wp.array2d(dtype=wp.vec2),
-    geom_type: wp.array(dtype=int),
+    geom_plugin_index: wp.array[int],
+    geom_priority: wp.array[int],
+    geom_rbound: wp.array2d[float],
+    geom_size: wp.array2d[wp.vec3],
+    geom_solimp: wp.array2d[mjwp_types.vec5],
+    geom_solmix: wp.array2d[float],
+    geom_solref: wp.array2d[wp.vec2],
+    geom_type: wp.array[int],
     has_sdf_geom: bool,
-    hfield_adr: wp.array(dtype=int),
-    hfield_data: wp.array(dtype=float),
-    hfield_ncol: wp.array(dtype=int),
-    hfield_nrow: wp.array(dtype=int),
-    hfield_size: wp.array(dtype=wp.vec4),
-    mesh_face: wp.array(dtype=wp.vec3i),
-    mesh_faceadr: wp.array(dtype=int),
-    mesh_graph: wp.array(dtype=int),
-    mesh_graphadr: wp.array(dtype=int),
-    mesh_octadr: wp.array(dtype=int),
-    mesh_polyadr: wp.array(dtype=int),
-    mesh_polymap: wp.array(dtype=int),
-    mesh_polymapadr: wp.array(dtype=int),
-    mesh_polymapnum: wp.array(dtype=int),
-    mesh_polynormal: wp.array(dtype=wp.vec3),
-    mesh_polynum: wp.array(dtype=int),
-    mesh_polyvert: wp.array(dtype=int),
-    mesh_polyvertadr: wp.array(dtype=int),
-    mesh_polyvertnum: wp.array(dtype=int),
-    mesh_vert: wp.array(dtype=wp.vec3),
-    mesh_vertadr: wp.array(dtype=int),
-    mesh_vertnum: wp.array(dtype=int),
+    hfield_adr: wp.array[int],
+    hfield_data: wp.array[float],
+    hfield_ncol: wp.array[int],
+    hfield_nrow: wp.array[int],
+    hfield_size: wp.array[wp.vec4],
+    mesh_face: wp.array[wp.vec3i],
+    mesh_faceadr: wp.array[int],
+    mesh_graph: wp.array[int],
+    mesh_graphadr: wp.array[int],
+    mesh_octadr: wp.array[int],
+    mesh_polyadr: wp.array[int],
+    mesh_polymap: wp.array[int],
+    mesh_polymapadr: wp.array[int],
+    mesh_polymapnum: wp.array[int],
+    mesh_polynormal: wp.array[wp.vec3],
+    mesh_polynum: wp.array[int],
+    mesh_polyvert: wp.array[int],
+    mesh_polyvertadr: wp.array[int],
+    mesh_polyvertnum: wp.array[int],
+    mesh_vert: wp.array[wp.vec3],
+    mesh_vertadr: wp.array[int],
+    mesh_vertnum: wp.array[int],
     nflex: int,
     nflexelem: int,
     nflexshelldata: int,
@@ -115,25 +116,25 @@ def _collision_shim(
     nmaxmeshdeg: int,
     nmaxpolygon: int,
     nmeshface: int,
-    nxn_geom_pair_filtered: wp.array(dtype=wp.vec2i),
-    nxn_pairid: wp.array(dtype=wp.vec2i),
-    nxn_pairid_filtered: wp.array(dtype=wp.vec2i),
-    oct_aabb: wp.array2d(dtype=wp.vec3),
-    oct_child: wp.array(dtype=mjwp_types.vec8i),
-    oct_coeff: wp.array(dtype=mjwp_types.vec8),
-    pair_dim: wp.array(dtype=int),
-    pair_friction: wp.array2d(dtype=mjwp_types.vec5),
-    pair_gap: wp.array2d(dtype=float),
-    pair_margin: wp.array2d(dtype=float),
-    pair_solimp: wp.array2d(dtype=mjwp_types.vec5),
-    pair_solref: wp.array2d(dtype=wp.vec2),
-    pair_solreffriction: wp.array2d(dtype=wp.vec2),
-    plugin: wp.array(dtype=int),
-    plugin_attr: wp.array(dtype=mjwp_types.vec_pluginattr),
+    nxn_geom_pair_filtered: wp.array[wp.vec2i],
+    nxn_pairid: wp.array[wp.vec2i],
+    nxn_pairid_filtered: wp.array[wp.vec2i],
+    oct_aabb: wp.array2d[wp.vec3],
+    oct_child: wp.array[mjwp_types.vec8i],
+    oct_coeff: wp.array[mjwp_types.vec8],
+    pair_dim: wp.array[int],
+    pair_friction: wp.array2d[mjwp_types.vec5],
+    pair_gap: wp.array2d[float],
+    pair_margin: wp.array2d[float],
+    pair_solimp: wp.array2d[mjwp_types.vec5],
+    pair_solref: wp.array2d[wp.vec2],
+    pair_solreffriction: wp.array2d[wp.vec2],
+    plugin: wp.array[int],
+    plugin_attr: wp.array[mjwp_types.vec_pluginattr],
     opt__broadphase: int,
     opt__broadphase_filter: int,
     opt__ccd_iterations: int,
-    opt__ccd_tolerance: wp.array(dtype=float),
+    opt__ccd_tolerance: wp.array[float],
     opt__disableflags: int,
     opt__enableflags: int,
     opt__sdf_initpoints: int,
@@ -141,27 +142,27 @@ def _collision_shim(
     # Data
     naccdmax: int,
     naconmax: int,
-    flexvert_xpos: wp.array2d(dtype=wp.vec3),
-    geom_xmat: wp.array2d(dtype=wp.mat33),
-    geom_xpos: wp.array2d(dtype=wp.vec3),
-    nacon: wp.array(dtype=int),
-    ncollision: wp.array(dtype=int),
-    contact__dim: wp.array(dtype=int),
-    contact__dist: wp.array(dtype=float),
-    contact__efc_address: wp.array2d(dtype=int),
-    contact__flex: wp.array(dtype=wp.vec2i),
-    contact__frame: wp.array(dtype=wp.mat33),
-    contact__friction: wp.array(dtype=mjwp_types.vec5),
-    contact__geom: wp.array(dtype=wp.vec2i),
-    contact__geomcollisionid: wp.array(dtype=int),
-    contact__includemargin: wp.array(dtype=float),
-    contact__pos: wp.array(dtype=wp.vec3),
-    contact__solimp: wp.array(dtype=mjwp_types.vec5),
-    contact__solref: wp.array(dtype=wp.vec2),
-    contact__solreffriction: wp.array(dtype=wp.vec2),
-    contact__type: wp.array(dtype=int),
-    contact__vert: wp.array(dtype=wp.vec2i),
-    contact__worldid: wp.array(dtype=int),
+    flexvert_xpos: wp.array2d[wp.vec3],
+    geom_xmat: wp.array2d[wp.mat33],
+    geom_xpos: wp.array2d[wp.vec3],
+    nacon: wp.array[int],
+    ncollision: wp.array[int],
+    contact__dim: wp.array[int],
+    contact__dist: wp.array[float],
+    contact__efc_address: wp.array2d[int],
+    contact__flex: wp.array[wp.vec2i],
+    contact__frame: wp.array[wp.mat33],
+    contact__friction: wp.array[mjwp_types.vec5],
+    contact__geom: wp.array[wp.vec2i],
+    contact__geomcollisionid: wp.array[int],
+    contact__includemargin: wp.array[float],
+    contact__pos: wp.array[wp.vec3],
+    contact__solimp: wp.array[mjwp_types.vec5],
+    contact__solref: wp.array[wp.vec2],
+    contact__solreffriction: wp.array[wp.vec2],
+    contact__type: wp.array[int],
+    contact__vert: wp.array[wp.vec2i],
+    contact__worldid: wp.array[int],
 ):
   _m.stat = _s
   _m.opt = _o
@@ -376,7 +377,7 @@ def _collision_jax_impl(m: types.Model, d: types.Data):
       m.geom_conaffinity,
       m.geom_condim,
       m.geom_contype,
-      m.geom_dataid,
+      jax.numpy.expand_dims(m.geom_dataid, 0),
       m.geom_friction,
       m.geom_gap,
       m.geom_margin,

mjx/mujoco/mjx/warp/render.py

@@ -48,42 +48,43 @@ _cb = mjwp_types.Callback(
     **{f.name: None for f in dataclasses.fields(mjwp_types.Callback) if f.init}
 )
 
+
 @ffi.format_args_for_warp
 def _render_shim(
     # Model
     nworld: int,
-    cam_fovy: wp.array2d(dtype=float),
-    cam_intrinsic: wp.array2d(dtype=wp.vec4),
-    cam_projection: wp.array(dtype=int),
-    cam_sensorsize: wp.array(dtype=wp.vec2),
-    flex_edge: wp.array(dtype=wp.vec2i),
-    flex_radius: wp.array(dtype=float),
-    flex_vertadr: wp.array(dtype=int),
-    geom_dataid: wp.array(dtype=int),
-    geom_matid: wp.array2d(dtype=int),
-    geom_rgba: wp.array2d(dtype=wp.vec4),
-    geom_size: wp.array2d(dtype=wp.vec3),
-    geom_type: wp.array(dtype=int),
-    light_active: wp.array2d(dtype=bool),
-    light_castshadow: wp.array2d(dtype=bool),
-    light_type: wp.array2d(dtype=int),
-    mat_rgba: wp.array2d(dtype=wp.vec4),
-    mat_texid: wp.array3d(dtype=int),
-    mat_texrepeat: wp.array2d(dtype=wp.vec2),
-    mesh_faceadr: wp.array(dtype=int),
+    cam_fovy: wp.array2d[float],
+    cam_intrinsic: wp.array2d[wp.vec4],
+    cam_projection: wp.array[int],
+    cam_sensorsize: wp.array[wp.vec2],
+    flex_edge: wp.array[wp.vec2i],
+    flex_radius: wp.array[float],
+    flex_vertadr: wp.array[int],
+    geom_dataid: wp.array2d[int],
+    geom_matid: wp.array2d[int],
+    geom_rgba: wp.array2d[wp.vec4],
+    geom_size: wp.array2d[wp.vec3],
+    geom_type: wp.array[int],
+    light_active: wp.array2d[bool],
+    light_castshadow: wp.array2d[bool],
+    light_type: wp.array2d[int],
+    mat_rgba: wp.array2d[wp.vec4],
+    mat_texid: wp.array3d[int],
+    mat_texrepeat: wp.array2d[wp.vec2],
+    mesh_faceadr: wp.array[int],
     nlight: int,
     # Data
-    cam_xmat: wp.array2d(dtype=wp.mat33),
-    cam_xpos: wp.array2d(dtype=wp.vec3),
-    flexvert_xpos: wp.array2d(dtype=wp.vec3),
-    geom_xmat: wp.array2d(dtype=wp.mat33),
-    geom_xpos: wp.array2d(dtype=wp.vec3),
-    light_xdir: wp.array2d(dtype=wp.vec3),
-    light_xpos: wp.array2d(dtype=wp.vec3),
+    cam_xmat: wp.array2d[wp.mat33],
+    cam_xpos: wp.array2d[wp.vec3],
+    flexvert_xpos: wp.array2d[wp.vec3],
+    geom_xmat: wp.array2d[wp.mat33],
+    geom_xpos: wp.array2d[wp.vec3],
+    light_xdir: wp.array2d[wp.vec3],
+    light_xpos: wp.array2d[wp.vec3],
     # Registry
     rc_id: int,
-    rgb: wp.array2d(dtype=wp.uint32),
-    depth: wp.array2d(dtype=wp.float32),
+    rgb: wp.array2d[wp.uint32],
+    depth: wp.array2d[wp.float32],
 ):
   _m.stat = _s
   _m.opt = _o
@@ -164,7 +165,7 @@ def _render_jax_impl(m: types.Model, d: types.Data, ctx: RenderContextPytree):
       m._impl.flex_edge,
       m._impl.flex_radius,
       m.flex_vertadr,
-      m.geom_dataid,
+      jax.numpy.expand_dims(m.geom_dataid, 0),
       m.geom_matid,
       m.geom_rgba,
       m.geom_size,

mjx/mujoco/mjx/warp/smooth.py

@@ -46,52 +46,53 @@ _cb = mjwp_types.Callback(
     **{f.name: None for f in dataclasses.fields(mjwp_types.Callback) if f.init}
 )
 
+
 @ffi.format_args_for_warp
 def _kinematics_shim(
     # Model
     nworld: int,
-    body_branch_start: wp.array(dtype=int),
-    body_branches: wp.array(dtype=int),
-    body_ipos: wp.array2d(dtype=wp.vec3),
-    body_iquat: wp.array2d(dtype=wp.quat),
-    body_jntadr: wp.array(dtype=int),
-    body_jntnum: wp.array(dtype=int),
-    body_mocapid: wp.array(dtype=int),
-    body_parentid: wp.array(dtype=int),
-    body_pos: wp.array2d(dtype=wp.vec3),
-    body_quat: wp.array2d(dtype=wp.quat),
-    body_rootid: wp.array(dtype=int),
-    body_weldid: wp.array(dtype=int),
-    geom_bodyid: wp.array(dtype=int),
-    geom_pos: wp.array2d(dtype=wp.vec3),
-    geom_quat: wp.array2d(dtype=wp.quat),
-    jnt_axis: wp.array2d(dtype=wp.vec3),
-    jnt_pos: wp.array2d(dtype=wp.vec3),
-    jnt_qposadr: wp.array(dtype=int),
-    jnt_type: wp.array(dtype=int),
+    body_branch_start: wp.array[int],
+    body_branches: wp.array[int],
+    body_ipos: wp.array2d[wp.vec3],
+    body_iquat: wp.array2d[wp.quat],
+    body_jntadr: wp.array[int],
+    body_jntnum: wp.array[int],
+    body_mocapid: wp.array[int],
+    body_parentid: wp.array[int],
+    body_pos: wp.array2d[wp.vec3],
+    body_quat: wp.array2d[wp.quat],
+    body_rootid: wp.array[int],
+    body_weldid: wp.array[int],
+    geom_bodyid: wp.array[int],
+    geom_pos: wp.array2d[wp.vec3],
+    geom_quat: wp.array2d[wp.quat],
+    jnt_axis: wp.array2d[wp.vec3],
+    jnt_pos: wp.array2d[wp.vec3],
+    jnt_qposadr: wp.array[int],
+    jnt_type: wp.array[int],
     nbody: int,
     nbranch: int,
     ngeom: int,
     nsite: int,
-    qpos0: wp.array2d(dtype=float),
-    site_bodyid: wp.array(dtype=int),
-    site_pos: wp.array2d(dtype=wp.vec3),
-    site_quat: wp.array2d(dtype=wp.quat),
+    qpos0: wp.array2d[float],
+    site_bodyid: wp.array[int],
+    site_pos: wp.array2d[wp.vec3],
+    site_quat: wp.array2d[wp.quat],
     # Data
-    geom_xmat: wp.array2d(dtype=wp.mat33),
-    geom_xpos: wp.array2d(dtype=wp.vec3),
-    mocap_pos: wp.array2d(dtype=wp.vec3),
-    mocap_quat: wp.array2d(dtype=wp.quat),
-    qpos: wp.array2d(dtype=float),
-    site_xmat: wp.array2d(dtype=wp.mat33),
-    site_xpos: wp.array2d(dtype=wp.vec3),
-    xanchor: wp.array2d(dtype=wp.vec3),
-    xaxis: wp.array2d(dtype=wp.vec3),
-    ximat: wp.array2d(dtype=wp.mat33),
-    xipos: wp.array2d(dtype=wp.vec3),
-    xmat: wp.array2d(dtype=wp.mat33),
-    xpos: wp.array2d(dtype=wp.vec3),
-    xquat: wp.array2d(dtype=wp.quat),
+    geom_xmat: wp.array2d[wp.mat33],
+    geom_xpos: wp.array2d[wp.vec3],
+    mocap_pos: wp.array2d[wp.vec3],
+    mocap_quat: wp.array2d[wp.quat],
+    qpos: wp.array2d[float],
+    site_xmat: wp.array2d[wp.mat33],
+    site_xpos: wp.array2d[wp.vec3],
+    xanchor: wp.array2d[wp.vec3],
+    xaxis: wp.array2d[wp.vec3],
+    ximat: wp.array2d[wp.mat33],
+    xipos: wp.array2d[wp.vec3],
+    xmat: wp.array2d[wp.mat33],
+    xpos: wp.array2d[wp.vec3],
+    xquat: wp.array2d[wp.quat],
 ):
   _m.stat = _s
   _m.opt = _o
@@ -297,45 +298,45 @@ def kinematics_vmap(
 def _tendon_shim(
     # Model
     nworld: int,
-    body_dofadr: wp.array(dtype=int),
-    body_dofnum: wp.array(dtype=int),
-    body_parentid: wp.array(dtype=int),
-    body_rootid: wp.array(dtype=int),
-    geom_bodyid: wp.array(dtype=int),
-    geom_size: wp.array2d(dtype=wp.vec3),
-    jnt_dofadr: wp.array(dtype=int),
-    jnt_qposadr: wp.array(dtype=int),
+    body_dofadr: wp.array[int],
+    body_dofnum: wp.array[int],
+    body_parentid: wp.array[int],
+    body_rootid: wp.array[int],
+    geom_bodyid: wp.array[int],
+    geom_size: wp.array2d[wp.vec3],
+    jnt_dofadr: wp.array[int],
+    jnt_qposadr: wp.array[int],
     ntendon: int,
     nwrap: int,
-    site_bodyid: wp.array(dtype=int),
-    ten_J_colind: wp.array(dtype=int),
-    ten_J_rowadr: wp.array(dtype=int),
-    ten_J_rownnz: wp.array(dtype=int),
-    tendon_adr: wp.array(dtype=int),
-    tendon_geom_adr: wp.array(dtype=int),
-    tendon_jnt_adr: wp.array(dtype=int),
-    tendon_num: wp.array(dtype=int),
-    tendon_site_pair_adr: wp.array(dtype=int),
-    wrap_geom_adr: wp.array(dtype=int),
-    wrap_jnt_adr: wp.array(dtype=int),
-    wrap_objid: wp.array(dtype=int),
-    wrap_prm: wp.array(dtype=float),
-    wrap_pulley_scale: wp.array(dtype=float),
-    wrap_site_pair_adr: wp.array(dtype=int),
-    wrap_type: wp.array(dtype=int),
+    site_bodyid: wp.array[int],
+    ten_J_colind: wp.array[int],
+    ten_J_rowadr: wp.array[int],
+    ten_J_rownnz: wp.array[int],
+    tendon_adr: wp.array[int],
+    tendon_geom_adr: wp.array[int],
+    tendon_jnt_adr: wp.array[int],
+    tendon_num: wp.array[int],
+    tendon_site_pair_adr: wp.array[int],
+    wrap_geom_adr: wp.array[int],
+    wrap_jnt_adr: wp.array[int],
+    wrap_objid: wp.array[int],
+    wrap_prm: wp.array[float],
+    wrap_pulley_scale: wp.array[float],
+    wrap_site_pair_adr: wp.array[int],
+    wrap_type: wp.array[int],
     # Data
-    cdof: wp.array2d(dtype=wp.spatial_vector),
-    geom_xmat: wp.array2d(dtype=wp.mat33),
-    geom_xpos: wp.array2d(dtype=wp.vec3),
-    qpos: wp.array2d(dtype=float),
-    site_xpos: wp.array2d(dtype=wp.vec3),
-    subtree_com: wp.array2d(dtype=wp.vec3),
-    ten_J: wp.array2d(dtype=float),
-    ten_length: wp.array2d(dtype=float),
-    ten_wrapadr: wp.array2d(dtype=int),
-    ten_wrapnum: wp.array2d(dtype=int),
-    wrap_obj: wp.array2d(dtype=wp.vec2i),
-    wrap_xpos: wp.array2d(dtype=wp.spatial_vector),
+    cdof: wp.array2d[wp.spatial_vector],
+    geom_xmat: wp.array2d[wp.mat33],
+    geom_xpos: wp.array2d[wp.vec3],
+    qpos: wp.array2d[float],
+    site_xpos: wp.array2d[wp.vec3],
+    subtree_com: wp.array2d[wp.vec3],
+    ten_J: wp.array2d[float],
+    ten_length: wp.array2d[float],
+    ten_wrapadr: wp.array2d[int],
+    ten_wrapnum: wp.array2d[int],
+    wrap_obj: wp.array2d[wp.vec2i],
+    wrap_xpos: wp.array2d[wp.spatial_vector],
 ):
   _m.stat = _s
   _m.opt = _o
@@ -489,26 +490,26 @@ def tendon_vmap(unused_axis_size, is_batched, m: types.Model, d: types.Data):
 def _com_pos_shim(
     # Model
     nworld: int,
-    body_inertia: wp.array2d(dtype=wp.vec3),
-    body_mass: wp.array2d(dtype=float),
-    body_parentid: wp.array(dtype=int),
-    body_rootid: wp.array(dtype=int),
-    body_subtreemass: wp.array2d(dtype=float),
-    body_tree: tuple[wp.array(dtype=int), ...],
-    jnt_bodyid: wp.array(dtype=int),
-    jnt_dofadr: wp.array(dtype=int),
-    jnt_type: wp.array(dtype=int),
+    body_inertia: wp.array2d[wp.vec3],
+    body_mass: wp.array2d[float],
+    body_parentid: wp.array[int],
+    body_rootid: wp.array[int],
+    body_subtreemass: wp.array2d[float],
+    body_tree: tuple[wp.array[int], ...],
+    jnt_bodyid: wp.array[int],
+    jnt_dofadr: wp.array[int],
+    jnt_type: wp.array[int],
     nbody: int,
     njnt: int,
     # Data
-    cdof: wp.array2d(dtype=wp.spatial_vector),
-    cinert: wp.array2d(dtype=mjwp_types.vec10),
-    subtree_com: wp.array2d(dtype=wp.vec3),
-    xanchor: wp.array2d(dtype=wp.vec3),
-    xaxis: wp.array2d(dtype=wp.vec3),
-    ximat: wp.array2d(dtype=wp.mat33),
-    xipos: wp.array2d(dtype=wp.vec3),
-    xmat: wp.array2d(dtype=wp.mat33),
+    cdof: wp.array2d[wp.spatial_vector],
+    cinert: wp.array2d[mjwp_types.vec10],
+    subtree_com: wp.array2d[wp.vec3],
+    xanchor: wp.array2d[wp.vec3],
+    xaxis: wp.array2d[wp.vec3],
+    ximat: wp.array2d[wp.mat33],
+    xipos: wp.array2d[wp.vec3],
+    xmat: wp.array2d[wp.mat33],
 ):
   _m.stat = _s
   _m.opt = _o

mjx/mujoco/mjx/warp/types.py

@@ -672,7 +672,7 @@ _NDIM = {
         'geom_conaffinity': 1,
         'geom_condim': 1,
         'geom_contype': 1,
-        'geom_dataid': 1,
+        'geom_dataid': 2,
         'geom_fluid': 2,
         'geom_friction': 3,
         'geom_gap': 2,
@@ -1238,7 +1238,7 @@ _BATCH_DIM = {
         'geom_conaffinity': False,
         'geom_condim': False,
         'geom_contype': False,
-        'geom_dataid': False,
+        'geom_dataid': True,
         'geom_fluid': False,
         'geom_friction': True,
         'geom_gap': True,