diff --git a/doc/changelog.rst b/doc/changelog.rst
index ee10dc2a..7d8c9bb9 100644
--- a/doc/changelog.rst
+++ b/doc/changelog.rst
@@ -2,6 +2,19 @@
 Changelog
 =========
 
+Upcoming version (not yet released)
+-----------------------------------
+
+General
+^^^^^^^
+
+MJX
+^^^
+1. Added ``actuator_length``, ``cdof`` and ``cdof_dof`` fields to ``mjx.Data``.
+
+Bug fixes
+^^^^^^^^^
+
 Version 3.4.0 (December 5, 2025)
 --------------------------------
 
diff --git a/mjx/mujoco/mjx/_src/forward.py b/mjx/mujoco/mjx/_src/forward.py
index 2a0e946e..088728ff 100644
--- a/mjx/mujoco/mjx/_src/forward.py
+++ b/mjx/mujoco/mjx/_src/forward.py
@@ -182,7 +182,7 @@ def fwd_actuation(m: Model, d: Data) -> Data:
       m.actuator_gainprm,
       m.actuator_biastype,
       m.actuator_biasprm,
-      d._impl.actuator_length,
+      d.actuator_length,
       d._impl.actuator_velocity,
       ctrl_act,
       jp.array(m.actuator_lengthrange),
diff --git a/mjx/mujoco/mjx/_src/io.py b/mjx/mujoco/mjx/_src/io.py
index cec570a1..c1012dc0 100644
--- a/mjx/mujoco/mjx/_src/io.py
+++ b/mjx/mujoco/mjx/_src/io.py
@@ -552,6 +552,7 @@ def _make_data_public_fields(m: types.Model) -> Dict[str, Any]:
       'cam_xmat': (m.ncam, 3, 3, float_),
       'subtree_com': (m.nbody, 3, float_),
       'actuator_force': (m.nu, float_),
+      'actuator_length': (m.nu, float_),
       'qfrc_bias': (m.nv, float_),
       'qfrc_gravcomp': (m.nv, float_),
       'qfrc_fluid': (m.nv, float_),
@@ -562,6 +563,8 @@ def _make_data_public_fields(m: types.Model) -> Dict[str, Any]:
       'qfrc_constraint': (m.nv, float_),
       'qfrc_inverse': (m.nv, float_),
       'cvel': (m.nbody, 6, float_),
+      'cdof': (m.nv, 6, float_),
+      'cdof_dot': (m.nv, 6, float_),
       'ten_length': (m.ntendon, float_),
   }
   zero_fields = {
@@ -619,14 +622,12 @@ def _make_data_jax(
 
   zero_impl_fields = {
       'solver_niter': (int_,),
-      'cdof': (m.nv, 6, float_),
       'cinert': (m.nbody, 10, float_),
       'ten_wrapadr': (m.ntendon, np.int32),
       'ten_wrapnum': (m.ntendon, np.int32),
       'ten_J': (m.ntendon, m.nv, float_),
       'wrap_obj': (m.nwrap, 2, np.int32),
       'wrap_xpos': (m.nwrap, 6, float_),
-      'actuator_length': (m.nu, float_),
       'actuator_moment': (m.nu, m.nv, float_),
       'crb': (m.nbody, 10, float_),
       'qM': (m.nM, float_) if support.is_sparse(m) else (m.nv, m.nv, float_),
@@ -635,7 +636,6 @@ def _make_data_jax(
       'qLDiagInv': (m.nv, float_) if support.is_sparse(m) else (0, float_),
       'ten_velocity': (m.ntendon, float_),
       'actuator_velocity': (m.nu, float_),
-      'cdof_dot': (m.nv, 6, float_),
       'cacc': (m.nbody, 6, float_),
       'cfrc_int': (m.nbody, 6, float_),
       'cfrc_ext': (m.nbody, 6, float_),
@@ -713,7 +713,6 @@ def _make_data_c(
   nbvhdynamic = get(m, 'nbvhdynamic')
   zero_impl_fields = {
       'solver_niter': (int_,),
-      'cdof': (m.nv, 6, float_),
       'cinert': (m.nbody, 10, float_),
       'light_xpos': (m.nlight, 3, float_),
       'light_xdir': (m.nlight, 3, float_),
@@ -732,7 +731,6 @@ def _make_data_c(
       'ten_wrapnum': (m.ntendon, np.int32),
       'wrap_obj': (m.nwrap, 2, np.int32),
       'wrap_xpos': (m.nwrap, 6, float_),
-      'actuator_length': (m.nu, float_),
       'moment_rownnz': (m.nu, np.int32),
       'moment_rowadr': (m.nu, np.int32),
       'moment_colind': (m.nJmom, np.int32),
@@ -762,7 +760,6 @@ def _make_data_c(
       'qLU': (m.nD, float_),
       'qfrc_spring': (m.nv, float_),
       'qfrc_damper': (m.nv, float_),
-      'cdof_dot': (m.nv, 6, float_),
       'cacc': (m.nbody, 6, float_),
       'cfrc_int': (m.nbody, 6, float_),
       'cfrc_ext': (m.nbody, 6, float_),
diff --git a/mjx/mujoco/mjx/_src/io_test.py b/mjx/mujoco/mjx/_src/io_test.py
index c9eb7cb8..63d74bdf 100644
--- a/mjx/mujoco/mjx/_src/io_test.py
+++ b/mjx/mujoco/mjx/_src/io_test.py
@@ -402,10 +402,10 @@ class DataIOTest(parameterized.TestCase):
     self.assertEqual(d.geom_xpos.shape, (6, 3))
     self.assertEqual(d.geom_xmat.shape, (6, 3, 3))
     self.assertEqual(d.subtree_com.shape, (nbody, 3))
-    self.assertEqual(d._impl.cdof.shape, (nv, 6))
+    self.assertEqual(d.cdof.shape, (nv, 6))
     self.assertEqual(d._impl.cinert.shape, (nbody, 10))
     self.assertEqual(d._impl.crb.shape, (nbody, 10))
-    self.assertEqual(d._impl.actuator_length.shape, (1,))
+    self.assertEqual(d.actuator_length.shape, (1,))
     if impl == 'jax':
       self.assertEqual(d._impl.actuator_moment.shape, (1, nv))
     elif impl == 'c':
@@ -422,7 +422,7 @@ class DataIOTest(parameterized.TestCase):
     self.assertEqual(d._impl.efc_D.shape, (nefc,))
     self.assertEqual(d._impl.actuator_velocity.shape, (1,))
     self.assertEqual(d.cvel.shape, (nbody, 6))
-    self.assertEqual(d._impl.cdof_dot.shape, (nv, 6))
+    self.assertEqual(d.cdof_dot.shape, (nv, 6))
     self.assertEqual(d.qfrc_bias.shape, (nv,))
     self.assertEqual(d.qfrc_passive.shape, (nv,))
     self.assertEqual(d._impl.efc_aref.shape, (nefc,))
@@ -477,7 +477,7 @@ class DataIOTest(parameterized.TestCase):
     np.testing.assert_allclose(dx.qpos, d.qpos)
     np.testing.assert_allclose(dx.xpos, d.xpos)
     np.testing.assert_allclose(dx.cvel, d.cvel)
-    np.testing.assert_allclose(dx._impl.cdof_dot, d.cdof_dot)
+    np.testing.assert_allclose(dx.cdof_dot, d.cdof_dot)
 
     # check that there are no weak types
     self.assertFalse(
diff --git a/mjx/mujoco/mjx/_src/sensor.py b/mjx/mujoco/mjx/_src/sensor.py
index 7d7fc6ed..7e0930cd 100644
--- a/mjx/mujoco/mjx/_src/sensor.py
+++ b/mjx/mujoco/mjx/_src/sensor.py
@@ -175,7 +175,7 @@ def sensor_pos(m: Model, d: Data) -> Data:
     elif sensor_type == SensorType.TENDONPOS:
       sensor = d.ten_length[objid]
     elif sensor_type == SensorType.ACTUATORPOS:
-      sensor = d._impl.actuator_length[objid]
+      sensor = d.actuator_length[objid]
     elif sensor_type == SensorType.BALLQUAT:
       jnt_qposadr = m.jnt_qposadr[objid, None] + np.arange(4)[None]
       quat = d.qpos[jnt_qposadr]
diff --git a/mjx/mujoco/mjx/_src/smooth.py b/mjx/mujoco/mjx/_src/smooth.py
index dd6aa89e..ab8fcea6 100644
--- a/mjx/mujoco/mjx/_src/smooth.py
+++ b/mjx/mujoco/mjx/_src/smooth.py
@@ -212,7 +212,7 @@ def com_pos(m: Model, d: Data) -> Data:
       d.xanchor,
       d.xaxis,
   )
-  d = d.tree_replace({'_impl.cdof': cdof})
+  d = d.tree_replace({'cdof': cdof})
 
   return d
 
@@ -305,8 +305,8 @@ def crb(m: Model, d: Data) -> Data:
   d = d.tree_replace({'_impl.crb': crb_body})
 
   crb_dof = jp.take(crb_body, jp.array(m.dof_bodyid), axis=0)
-  crb_cdof = jax.vmap(math.inert_mul)(crb_dof, d._impl.cdof)
-  qm = support.make_m(m, crb_cdof, d._impl.cdof, m.dof_armature)
+  crb_cdof = jax.vmap(math.inert_mul)(crb_dof, d.cdof)
+  qm = support.make_m(m, crb_cdof, d.cdof, m.dof_armature)
   d = d.tree_replace({'_impl.qM': qm})
   return d
 
@@ -445,11 +445,11 @@ def com_vel(m: Model, d: Data) -> Data:
       'jvv',
       'bv',
       m.jnt_type,
-      d._impl.cdof,
+      d.cdof,
       d.qvel,
   )
 
-  d = d.tree_replace({'cvel': cvel, '_impl.cdof_dot': cdof_dot})
+  d = d.tree_replace({'cvel': cvel, 'cdof_dot': cdof_dot})
 
   return d
 
@@ -576,7 +576,7 @@ def rne(m: Model, d: Data, flg_acc: bool = False) -> Data:
     return cacc
 
   cacc = scan.body_tree(
-      m, cacc_fn, 'vvvv', 'b', d._impl.cdof_dot, d.qvel, d._impl.cdof, d.qacc
+      m, cacc_fn, 'vvvv', 'b', d.cdof_dot, d.qvel, d.cdof, d.qacc
   )
 
   def frc(cinert, cacc, cvel):
@@ -594,7 +594,7 @@ def rne(m: Model, d: Data, flg_acc: bool = False) -> Data:
     return cfrc
 
   cfrc = scan.body_tree(m, cfrc_fn, 'b', 'b', loc_cfrc, reverse=True)
-  qfrc_bias = jax.vmap(jp.dot)(d._impl.cdof, cfrc[jp.array(m.dof_bodyid)])
+  qfrc_bias = jax.vmap(jp.dot)(d.cdof, cfrc[jp.array(m.dof_bodyid)])
 
   d = d.replace(qfrc_bias=qfrc_bias)
 
@@ -806,8 +806,8 @@ def rne_postconstraint(m: Model, d: Data) -> Data:
     )
 
     # cacc = cacc_parent + cdofdot * qvel + cdof * qacc
-    cacc_vel = d._impl.cdof_dot.T @ (mask * d.qvel)
-    cacc_acc = d._impl.cdof.T @ (mask * d.qacc)
+    cacc_vel = d.cdof_dot.T @ (mask * d.qvel)
+    cacc_acc = d.cdof.T @ (mask * d.qacc)
     cacc = cacc_parent + cacc_vel + cacc_acc
 
     # cfrc_body = cinert * cacc + cvel x (cinert * cvel)
@@ -1310,7 +1310,7 @@ def transmission(m: Model, d: Data) -> Data:
   moment = moment.reshape((m.nu, m.nv))
 
   d = d.tree_replace(
-      {'_impl.actuator_length': length, '_impl.actuator_moment': moment}
+      {'actuator_length': length, '_impl.actuator_moment': moment}
   )
   return d
 
diff --git a/mjx/mujoco/mjx/_src/smooth_test.py b/mjx/mujoco/mjx/_src/smooth_test.py
index 4e8b3711..302547da 100644
--- a/mjx/mujoco/mjx/_src/smooth_test.py
+++ b/mjx/mujoco/mjx/_src/smooth_test.py
@@ -81,7 +81,7 @@ class SmoothTest(absltest.TestCase):
     dx = jax.jit(mjx.com_pos)(mx, mjx.put_data(m, d))
     _assert_attr_eq(d, dx, 'subtree_com')
     _assert_attr_eq(d, dx._impl, 'cinert')
-    _assert_attr_eq(d, dx._impl, 'cdof')
+    _assert_attr_eq(d, dx, 'cdof')
     # camlight
     dx = jax.jit(mjx.camlight)(mx, mjx.put_data(m, d))
     _assert_attr_eq(d, dx, 'cam_xpos')
@@ -100,7 +100,7 @@ class SmoothTest(absltest.TestCase):
     # com_vel
     dx = jax.jit(mjx.com_vel)(mx, mjx.put_data(m, d))
     _assert_attr_eq(d, dx, 'cvel')
-    _assert_attr_eq(d, dx._impl, 'cdof_dot')
+    _assert_attr_eq(d, dx, 'cdof_dot')
     # rne
     dx = jax.jit(mjx.rne)(mx, mjx.put_data(m, d))
     _assert_attr_eq(d, dx, 'qfrc_bias')
@@ -125,7 +125,7 @@ class SmoothTest(absltest.TestCase):
     _assert_attr_eq(d, dx, 'ten_length')
     # transmission
     dx = jax.jit(mjx.transmission)(mx, dx)
-    _assert_attr_eq(d, dx._impl, 'actuator_length')
+    _assert_attr_eq(d, dx, 'actuator_length')
 
     # convert sparse actuator_moment to dense representation
     moment = np.zeros((m.nu, m.nv))
@@ -196,7 +196,7 @@ class SmoothTest(absltest.TestCase):
 
     mujoco.mj_transmission(m, d)
     dx = jax.jit(mjx.transmission)(mx, dx)
-    _assert_attr_eq(d, dx._impl, 'actuator_length')
+    _assert_attr_eq(d, dx, 'actuator_length')
 
     # convert sparse actuator_moment to dense representation
     moment = np.zeros((m.nu, m.nv))
diff --git a/mjx/mujoco/mjx/_src/support.py b/mjx/mujoco/mjx/_src/support.py
index 04b0a0d4..24031821 100644
--- a/mjx/mujoco/mjx/_src/support.py
+++ b/mjx/mujoco/mjx/_src/support.py
@@ -149,9 +149,9 @@ def jac(
   mask = mask[jp.array(m.dof_bodyid)] > 0
 
   offset = point - d.subtree_com[jp.array(m.body_rootid)[body_id]]
-  jacp = jax.vmap(lambda a, b=offset: a[3:] + jp.cross(a[:3], b))(d._impl.cdof)  # pytype: disable=attribute-error
+  jacp = jax.vmap(lambda a, b=offset: a[3:] + jp.cross(a[:3], b))(d.cdof)  # pytype: disable=attribute-error
   jacp = jax.vmap(jp.multiply)(jacp, mask)
-  jacr = jax.vmap(jp.multiply)(d._impl.cdof[:, :3], mask)  # pytype: disable=attribute-error
+  jacr = jax.vmap(jp.multiply)(d.cdof[:, :3], mask)  # pytype: disable=attribute-error
 
   return jacp, jacr
 
@@ -169,8 +169,8 @@ def jac_dot(
   offset = point - d.subtree_com[jp.array(m.body_rootid)[body_id]]
   pvel_lin = d.cvel[body_id][3:] - jp.cross(offset, d.cvel[body_id][:3])
 
-  cdof = d._impl.cdof
-  cdof_dot = d._impl.cdof_dot
+  cdof = d.cdof
+  cdof_dot = d.cdof_dot
 
   # check for quaternion
   jnt_type = m.jnt_type[m.dof_jntid]
diff --git a/mjx/mujoco/mjx/_src/types.py b/mjx/mujoco/mjx/_src/types.py
index 3560e32f..e13a0c5a 100644
--- a/mjx/mujoco/mjx/_src/types.py
+++ b/mjx/mujoco/mjx/_src/types.py
@@ -1021,7 +1021,6 @@ class DataC(PyTreeNode):
   plugin_data: jax.Array
   light_xpos: jax.Array
   light_xdir: jax.Array
-  cdof: jax.Array
   cinert: jax.Array
   flexvert_xpos: jax.Array
   flexelem_aabb: jax.Array
@@ -1039,7 +1038,6 @@ class DataC(PyTreeNode):
   ten_J: jax.Array  # pylint:disable=invalid-name
   wrap_obj: jax.Array
   wrap_xpos: jax.Array
-  actuator_length: jax.Array
   moment_rownnz: jax.Array  # pylint:disable=invalid-name
   moment_rowadr: jax.Array  # pylint:disable=invalid-name
   moment_colind: jax.Array  # pylint:disable=invalid-name
@@ -1059,7 +1057,7 @@ class DataC(PyTreeNode):
   flexedge_velocity: jax.Array
   ten_velocity: jax.Array
   actuator_velocity: jax.Array
-  cdof_dot: jax.Array
+
   qfrc_spring: jax.Array
   qfrc_damper: jax.Array
   subtree_linvel: jax.Array
@@ -1095,14 +1093,12 @@ class DataJAX(PyTreeNode):
   nefc: int
   ncon: int
   solver_niter: jax.Array
-  cdof: jax.Array
   cinert: jax.Array
   ten_wrapadr: jax.Array
   ten_wrapnum: jax.Array
   ten_J: jax.Array  # pylint:disable=invalid-name
   wrap_obj: jax.Array
   wrap_xpos: jax.Array
-  actuator_length: jax.Array
   actuator_moment: jax.Array
   crb: jax.Array
   qM: jax.Array  # pylint:disable=invalid-name
@@ -1111,7 +1107,7 @@ class DataJAX(PyTreeNode):
   qLDiagInv: jax.Array  # pylint:disable=invalid-name
   ten_velocity: jax.Array
   actuator_velocity: jax.Array
-  cdof_dot: jax.Array
+
   cacc: jax.Array
   cfrc_int: jax.Array
   cfrc_ext: jax.Array
@@ -1171,12 +1167,15 @@ class Data(PyTreeNode):
   cam_xmat: jax.Array
   subtree_com: jax.Array
   cvel: jax.Array
+  cdof: jax.Array
+  cdof_dot: jax.Array
   qfrc_bias: jax.Array
   qfrc_gravcomp: jax.Array
   qfrc_fluid: jax.Array
   qfrc_passive: jax.Array
   qfrc_actuator: jax.Array
   actuator_force: jax.Array
+  actuator_length: jax.Array
   qfrc_smooth: jax.Array
   qacc_smooth: jax.Array
   qfrc_constraint: jax.Array
diff --git a/mjx/mujoco/mjx/warp/collision_driver.py b/mjx/mujoco/mjx/warp/collision_driver.py
index 5e534e6b..446083ed 100644
--- a/mjx/mujoco/mjx/warp/collision_driver.py
+++ b/mjx/mujoco/mjx/warp/collision_driver.py
@@ -42,7 +42,6 @@ _e = mjwarp.Constraint(
     **{f.name: None for f in dataclasses.fields(mjwarp.Constraint) if f.init}
 )
 
-
 @ffi.format_args_for_warp
 def _collision_shim(
     # Model
diff --git a/mjx/mujoco/mjx/warp/forward.py b/mjx/mujoco/mjx/warp/forward.py
index 22d73d22..c5787fe1 100644
--- a/mjx/mujoco/mjx/warp/forward.py
+++ b/mjx/mujoco/mjx/warp/forward.py
@@ -42,7 +42,6 @@ _e = mjwarp.Constraint(
     **{f.name: None for f in dataclasses.fields(mjwarp.Constraint) if f.init}
 )
 
-
 @ffi.format_args_for_warp
 def _forward_shim(
     # Model
@@ -952,14 +951,14 @@ def _forward_jax_impl(m: types.Model, d: types.Data):
       'act': d.act.shape,
       'act_dot': d.act_dot.shape,
       'actuator_force': d.actuator_force.shape,
-      'actuator_length': d._impl.actuator_length.shape,
+      'actuator_length': d.actuator_length.shape,
       'actuator_moment': d._impl.actuator_moment.shape,
       'actuator_velocity': d._impl.actuator_velocity.shape,
       'cacc': d._impl.cacc.shape,
       'cam_xmat': d.cam_xmat.shape,
       'cam_xpos': d.cam_xpos.shape,
-      'cdof': d._impl.cdof.shape,
-      'cdof_dot': d._impl.cdof_dot.shape,
+      'cdof': d.cdof.shape,
+      'cdof_dot': d.cdof_dot.shape,
       'cfrc_ext': d._impl.cfrc_ext.shape,
       'cfrc_int': d._impl.cfrc_int.shape,
       'cinert': d._impl.cinert.shape,
@@ -1539,14 +1538,14 @@ def _forward_jax_impl(m: types.Model, d: types.Data):
       d.act,
       d.act_dot,
       d.actuator_force,
-      d._impl.actuator_length,
+      d.actuator_length,
       d._impl.actuator_moment,
       d._impl.actuator_velocity,
       d._impl.cacc,
       d.cam_xmat,
       d.cam_xpos,
-      d._impl.cdof,
-      d._impl.cdof_dot,
+      d.cdof,
+      d.cdof_dot,
       d._impl.cfrc_ext,
       d._impl.cfrc_int,
       d._impl.cinert,
@@ -1672,14 +1671,14 @@ def _forward_jax_impl(m: types.Model, d: types.Data):
       'act': out[0],
       'act_dot': out[1],
       'actuator_force': out[2],
-      '_impl.actuator_length': out[3],
+      'actuator_length': out[3],
       '_impl.actuator_moment': out[4],
       '_impl.actuator_velocity': out[5],
       '_impl.cacc': out[6],
       'cam_xmat': out[7],
       'cam_xpos': out[8],
-      '_impl.cdof': out[9],
-      '_impl.cdof_dot': out[10],
+      'cdof': out[9],
+      'cdof_dot': out[10],
       '_impl.cfrc_ext': out[11],
       '_impl.cfrc_int': out[12],
       '_impl.cinert': out[13],
@@ -1834,7 +1833,6 @@ _e = mjwarp.Constraint(
     **{f.name: None for f in dataclasses.fields(mjwarp.Constraint) if f.init}
 )
 
-
 @ffi.format_args_for_warp
 def _step_shim(
     # Model
@@ -2748,14 +2746,14 @@ def _step_jax_impl(m: types.Model, d: types.Data):
       'act': d.act.shape,
       'act_dot': d.act_dot.shape,
       'actuator_force': d.actuator_force.shape,
-      'actuator_length': d._impl.actuator_length.shape,
+      'actuator_length': d.actuator_length.shape,
       'actuator_moment': d._impl.actuator_moment.shape,
       'actuator_velocity': d._impl.actuator_velocity.shape,
       'cacc': d._impl.cacc.shape,
       'cam_xmat': d.cam_xmat.shape,
       'cam_xpos': d.cam_xpos.shape,
-      'cdof': d._impl.cdof.shape,
-      'cdof_dot': d._impl.cdof_dot.shape,
+      'cdof': d.cdof.shape,
+      'cdof_dot': d.cdof_dot.shape,
       'cfrc_ext': d._impl.cfrc_ext.shape,
       'cfrc_int': d._impl.cfrc_int.shape,
       'cinert': d._impl.cinert.shape,
@@ -3337,14 +3335,14 @@ def _step_jax_impl(m: types.Model, d: types.Data):
       d.act,
       d.act_dot,
       d.actuator_force,
-      d._impl.actuator_length,
+      d.actuator_length,
       d._impl.actuator_moment,
       d._impl.actuator_velocity,
       d._impl.cacc,
       d.cam_xmat,
       d.cam_xpos,
-      d._impl.cdof,
-      d._impl.cdof_dot,
+      d.cdof,
+      d.cdof_dot,
       d._impl.cfrc_ext,
       d._impl.cfrc_int,
       d._impl.cinert,
@@ -3470,14 +3468,14 @@ def _step_jax_impl(m: types.Model, d: types.Data):
       'act': out[0],
       'act_dot': out[1],
       'actuator_force': out[2],
-      '_impl.actuator_length': out[3],
+      'actuator_length': out[3],
       '_impl.actuator_moment': out[4],
       '_impl.actuator_velocity': out[5],
       '_impl.cacc': out[6],
       'cam_xmat': out[7],
       'cam_xpos': out[8],
-      '_impl.cdof': out[9],
-      '_impl.cdof_dot': out[10],
+      'cdof': out[9],
+      'cdof_dot': out[10],
       '_impl.cfrc_ext': out[11],
       '_impl.cfrc_int': out[12],
       '_impl.cinert': out[13],
diff --git a/mjx/mujoco/mjx/warp/forward_test.py b/mjx/mujoco/mjx/warp/forward_test.py
index 41376870..1412d34e 100644
--- a/mjx/mujoco/mjx/warp/forward_test.py
+++ b/mjx/mujoco/mjx/warp/forward_test.py
@@ -133,7 +133,7 @@ class ForwardTest(parameterized.TestCase):
       if m.nsite:
         tu.assert_attr_eq(dx, d, 'site_xpos')
         tu.assert_eq(dx.site_xmat, d.site_xmat.reshape((-1, 3, 3)), 'site_xmat')
-      tu.assert_attr_eq(dx._impl, d, 'cdof')
+      tu.assert_attr_eq(dx, d, 'cdof')
       tu.assert_attr_eq(dx._impl, d, 'cinert')
       tu.assert_attr_eq(dx, d, 'subtree_com')
       if m.nlight:
@@ -161,7 +161,7 @@ class ForwardTest(parameterized.TestCase):
 
       tu.assert_contact_eq(d, dx, worldid=i)
 
-      tu.assert_attr_eq(dx._impl, d, 'actuator_length')
+      tu.assert_attr_eq(dx, d, 'actuator_length')
       actuator_moment = np.zeros((m.nu, m.nv))
       mujoco.mju_sparse2dense(
           actuator_moment,
@@ -175,7 +175,7 @@ class ForwardTest(parameterized.TestCase):
       # fwd_velocity
       tu.assert_attr_eq(dx._impl, d, 'actuator_velocity')
       tu.assert_attr_eq(dx, d, 'cvel')
-      tu.assert_attr_eq(dx._impl, d, 'cdof_dot')
+      tu.assert_attr_eq(dx, d, 'cdof_dot')
       tu.assert_attr_eq(dx._impl, d, 'qfrc_spring')
       tu.assert_attr_eq(dx._impl, d, 'qfrc_damper')
       tu.assert_attr_eq(dx, d, 'qfrc_gravcomp')
diff --git a/mjx/mujoco/mjx/warp/smooth.py b/mjx/mujoco/mjx/warp/smooth.py
index 9a00ffb7..b1aaf5f5 100644
--- a/mjx/mujoco/mjx/warp/smooth.py
+++ b/mjx/mujoco/mjx/warp/smooth.py
@@ -42,7 +42,6 @@ _e = mjwarp.Constraint(
     **{f.name: None for f in dataclasses.fields(mjwarp.Constraint) if f.init}
 )
 
-
 @ffi.format_args_for_warp
 def _kinematics_shim(
     # Model
@@ -310,7 +309,6 @@ _e = mjwarp.Constraint(
     **{f.name: None for f in dataclasses.fields(mjwarp.Constraint) if f.init}
 )
 
-
 @ffi.format_args_for_warp
 def _tendon_shim(
     # Model
@@ -397,7 +395,7 @@ def _tendon_shim(
 
 def _tendon_jax_impl(m: types.Model, d: types.Data):
   output_dims = {
-      'cdof': d._impl.cdof.shape,
+      'cdof': d.cdof.shape,
       'geom_xmat': d.geom_xmat.shape,
       'geom_xpos': d.geom_xpos.shape,
       'qpos': d.qpos.shape,
@@ -455,7 +453,7 @@ def _tendon_jax_impl(m: types.Model, d: types.Data):
       m._impl.wrap_pulley_scale,
       m._impl.wrap_site_pair_adr,
       m.wrap_type,
-      d._impl.cdof,
+      d.cdof,
       d.geom_xmat,
       d.geom_xpos,
       d.qpos,
@@ -469,7 +467,7 @@ def _tendon_jax_impl(m: types.Model, d: types.Data):
       d._impl.wrap_xpos,
   )
   d = d.tree_replace({
-      '_impl.cdof': out[0],
+      'cdof': out[0],
       'geom_xmat': out[1],
       'geom_xpos': out[2],
       'qpos': out[3],
diff --git a/mjx/mujoco/mjx/warp/types.py b/mjx/mujoco/mjx/warp/types.py
index 71c3dcad..f13e6df2 100644
--- a/mjx/mujoco/mjx/warp/types.py
+++ b/mjx/mujoco/mjx/warp/types.py
@@ -217,12 +217,9 @@ class ModelWarp(PyTreeNode):
 
 class DataWarp(PyTreeNode):
   """Derived fields from Data."""
-  actuator_length: jax.Array
   actuator_moment: jax.Array
   actuator_velocity: jax.Array
   cacc: jax.Array
-  cdof: jax.Array
-  cdof_dot: jax.Array
   cfrc_ext: jax.Array
   cfrc_int: jax.Array
   cinert: jax.Array