From 05e26e961c3c9e61dceb4f0c0fe7ca94d8e4ffe0 Mon Sep 17 00:00:00 2001
From: Tom Power <tjpower@google.com>
Date: Thu, 9 Apr 2026 01:42:29 -0700
Subject: [PATCH] Refactor: Expose flex-related fields directly on mjx.Model.

PiperOrigin-RevId: 896958292
Change-Id: I94b32048e69d2b6844db0911f037318d5fbf4a6a
---
 mjx/mujoco/mjx/warp/bvh.py              |  7 +++----
 mjx/mujoco/mjx/warp/collision_driver.py |  5 ++---
 mjx/mujoco/mjx/warp/forward.py          | 13 ++++++-------
 mjx/mujoco/mjx/warp/render.py           |  3 +--
 mjx/mujoco/mjx/warp/types.py            |  3 ---
 5 files changed, 12 insertions(+), 19 deletions(-)

diff --git a/mjx/mujoco/mjx/warp/bvh.py b/mjx/mujoco/mjx/warp/bvh.py
index 2c1fee49..157315db 100644
--- a/mjx/mujoco/mjx/warp/bvh.py
+++ b/mjx/mujoco/mjx/warp/bvh.py
@@ -48,7 +48,6 @@ _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
@@ -134,12 +133,12 @@ def _refit_bvh_jax_impl(
       m._impl.flex_radius,
       m._impl.flex_shell,
       m._impl.flex_shelldataadr,
-      m._impl.flex_vertadr,
-      m._impl.flex_vertnum,
+      m.flex_vertadr,
+      m.flex_vertnum,
       m.geom_dataid,
       m.geom_size,
       m.geom_type,
-      m._impl.nflex,
+      m.nflex,
       m._impl.nflexelem,
       d._impl.flexvert_xpos,
       d.geom_xmat,
diff --git a/mjx/mujoco/mjx/warp/collision_driver.py b/mjx/mujoco/mjx/warp/collision_driver.py
index 9d5412d9..b0330c4b 100644
--- a/mjx/mujoco/mjx/warp/collision_driver.py
+++ b/mjx/mujoco/mjx/warp/collision_driver.py
@@ -46,7 +46,6 @@ _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
@@ -371,7 +370,7 @@ def _collision_jax_impl(m: types.Model, d: types.Data):
       m._impl.flex_shell,
       m._impl.flex_shelldataadr,
       m._impl.flex_shellnum,
-      m._impl.flex_vertadr,
+      m.flex_vertadr,
       m._impl.flex_vertflexid,
       m.geom_aabb,
       m.geom_conaffinity,
@@ -413,7 +412,7 @@ def _collision_jax_impl(m: types.Model, d: types.Data):
       m.mesh_vert,
       m.mesh_vertadr,
       m.mesh_vertnum,
-      m._impl.nflex,
+      m.nflex,
       m._impl.nflexelem,
       m._impl.nflexshelldata,
       m._impl.nflexvert,
diff --git a/mjx/mujoco/mjx/warp/forward.py b/mjx/mujoco/mjx/warp/forward.py
index e0f5c053..3d12681a 100644
--- a/mjx/mujoco/mjx/warp/forward.py
+++ b/mjx/mujoco/mjx/warp/forward.py
@@ -46,7 +46,6 @@ _cb = mjwp_types.Callback(
     **{f.name: None for f in dataclasses.fields(mjwp_types.Callback) if f.init}
 )
 
-
 @ffi.format_args_for_warp
 def _forward_shim(
     # Model
@@ -1499,10 +1498,10 @@ def _forward_jax_impl(m: types.Model, d: types.Data):
       m._impl.flex_shellnum,
       m._impl.flex_stiffness,
       m._impl.flex_vert,
-      m._impl.flex_vertadr,
+      m.flex_vertadr,
       m._impl.flex_vertbodyid,
       m._impl.flex_vertflexid,
-      m._impl.flex_vertnum,
+      m.flex_vertnum,
       m._impl.flexedge_J_colind,
       m._impl.flexedge_J_rowadr,
       m._impl.flexedge_J_rownnz,
@@ -1596,7 +1595,7 @@ def _forward_jax_impl(m: types.Model, d: types.Data):
       m._impl.nbranch,
       m.ncam,
       m.neq,
-      m._impl.nflex,
+      m.nflex,
       m._impl.nflexedge,
       m._impl.nflexelem,
       m._impl.nflexshelldata,
@@ -3444,10 +3443,10 @@ def _step_jax_impl(m: types.Model, d: types.Data):
       m._impl.flex_shellnum,
       m._impl.flex_stiffness,
       m._impl.flex_vert,
-      m._impl.flex_vertadr,
+      m.flex_vertadr,
       m._impl.flex_vertbodyid,
       m._impl.flex_vertflexid,
-      m._impl.flex_vertnum,
+      m.flex_vertnum,
       m._impl.flexedge_J_colind,
       m._impl.flexedge_J_rowadr,
       m._impl.flexedge_J_rownnz,
@@ -3542,7 +3541,7 @@ def _step_jax_impl(m: types.Model, d: types.Data):
       m._impl.nbranch,
       m.ncam,
       m.neq,
-      m._impl.nflex,
+      m.nflex,
       m._impl.nflexedge,
       m._impl.nflexelem,
       m._impl.nflexshelldata,
diff --git a/mjx/mujoco/mjx/warp/render.py b/mjx/mujoco/mjx/warp/render.py
index e723e6b0..6a3f66ab 100644
--- a/mjx/mujoco/mjx/warp/render.py
+++ b/mjx/mujoco/mjx/warp/render.py
@@ -48,7 +48,6 @@ _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
@@ -164,7 +163,7 @@ def _render_jax_impl(m: types.Model, d: types.Data, ctx: RenderContextPytree):
       m.cam_sensorsize,
       m._impl.flex_edge,
       m._impl.flex_radius,
-      m._impl.flex_vertadr,
+      m.flex_vertadr,
       m.geom_dataid,
       m.geom_matid,
       m.geom_rgba,
diff --git a/mjx/mujoco/mjx/warp/types.py b/mjx/mujoco/mjx/warp/types.py
index ef75edc4..564fe57e 100644
--- a/mjx/mujoco/mjx/warp/types.py
+++ b/mjx/mujoco/mjx/warp/types.py
@@ -165,10 +165,8 @@ class ModelWarp(PyTreeNode):
   flex_shellnum: np.ndarray
   flex_stiffness: np.ndarray
   flex_vert: np.ndarray
-  flex_vertadr: np.ndarray
   flex_vertbodyid: np.ndarray
   flex_vertflexid: np.ndarray
-  flex_vertnum: np.ndarray
   flexedge_J_colind: np.ndarray
   flexedge_J_rowadr: np.ndarray
   flexedge_J_rownnz: np.ndarray
@@ -200,7 +198,6 @@ class ModelWarp(PyTreeNode):
   nJfe: int
   nacttrnbody: int
   nbranch: int
-  nflex: int
   nflexedge: int
   nflexelem: int
   nflexelemdata: int