Rename backend_impl to impl. Remove _full_compat from callsites.

PiperOrigin-RevId: 770400560
Change-Id: I7bec79aeae07322a30b15cc04d4dd05077741a8b

mjx/mujoco/mjx/_src/io.py

@@ -41,51 +41,51 @@ def _is_cuda_gpu_device(device: jax.Device) -> bool:
   return device in cuda_devices
 
 
-def _resolve_backend_impl(
+def _resolve_impl(
     device: jax.Device,
-) -> types.BackendImpl:
-  """Pick a default backend impl based on the device specified."""
+) -> types.Impl:
+  """Pick a default implementation based on the device specified."""
   if _is_cuda_gpu_device(device):
     # TODO(btaba): Remove flag once Warp is ready to launch.
     mjx_warp_enabled = os.environ.get('MJX_WARP_ENABLED', 'f').lower() == 'true'
     if mjx_warp_enabled:
-      logging.debug('Picking default backend implementation: Warp.')
-      return types.BackendImpl.WARP
+      logging.debug('Picking default implementation: Warp.')
+      return types.Impl.WARP
     logging.info('MJX Warp is disabled via MJX_WARP_ENABLED=false.')
 
   if device.platform in ('gpu', 'tpu'):
-    logging.debug('Picking default backend implementation: JAX.')
-    return types.BackendImpl.JAX
+    logging.debug('Picking default implementation: JAX.')
+    return types.Impl.JAX
 
   if device.platform == 'cpu':
     mjx_c_default = (
         os.environ.get('MJX_C_DEFAULT_ENABLED', 'f').lower() == 'true'
     )
     if mjx_c_default:
-      logging.debug('Picking default backend implementation: C.')
-      return types.BackendImpl.C
-    return types.BackendImpl.JAX
+      logging.debug('Picking default implementation: C.')
+      return types.Impl.C
+    return types.Impl.JAX
 
   raise ValueError(f'Unsupported device: {device}')
 
 
 def _resolve_device(
-    backend_impl: types.BackendImpl,
+    impl: types.Impl,
 ) -> jax.Device:
-  """Resolves a device based on the backend implementation."""
-  backend_impl = types.BackendImpl(backend_impl)
-  if backend_impl == types.BackendImpl.JAX:
+  """Resolves a device based on the implementation."""
+  impl = types.Impl(impl)
+  if impl == types.Impl.JAX:
     device_0 = jax.devices()[0]
     logging.debug('Picking default device: %s.', device_0)
     return device_0
 
-  if backend_impl == types.BackendImpl.C:
+  if impl == types.Impl.C:
     cpu_0 = jax.devices('cpu')[0]
     logging.debug('Picking default device: %s', cpu_0)
     return cpu_0
 
-  if backend_impl == types.BackendImpl.WARP:
-    # WARP backend requires a CUDA GPU.
+  if impl == types.Impl.WARP:
+    # WARP implementation requires a CUDA GPU.
     cuda_gpus = [d for d in jax.devices('cuda')]
     if not cuda_gpus:
       raise AssertionError(
@@ -96,64 +96,64 @@ def _resolve_device(
     logging.debug('Picking default device: %s', cuda_gpus[0])
     return cuda_gpus[0]
 
-  raise ValueError(f'Unsupported backend implementation: {backend_impl}')
+  raise ValueError(f'Unsupported implementation: {impl}')
 
 
-def _check_backend_impl_device_compatibility(
-    backend_impl: Union[str, types.BackendImpl],
+def _check_impl_device_compatibility(
+    impl: Union[str, types.Impl],
     device: jax.Device,
 ) -> None:
-  """Checks that the backend implementation is compatible with the device."""
-  if backend_impl is None:
-    raise ValueError('No backend implementation specified.')
+  """Checks that the implementation is compatible with the device."""
+  if impl is None:
+    raise ValueError('No implementation specified.')
 
-  backend_impl = types.BackendImpl(backend_impl)
+  impl = types.Impl(impl)
 
-  if backend_impl == types.BackendImpl.WARP:
+  if impl == types.Impl.WARP:
     if not _is_cuda_gpu_device(device):
       raise AssertionError(
-          'Warp backend implementation requires a CUDA GPU device, got '
+          'Warp implementation requires a CUDA GPU device, got '
           f'{device}.'
       )
 
     mjx_warp_enabled = os.environ.get('MJX_WARP_ENABLED', 'f').lower() == 'true'
     if not mjx_warp_enabled:
       raise AssertionError(
-          'Warp backend implementation is disabled via MJX_WARP_ENABLED=false.'
+          'Warp implementation is disabled via MJX_WARP_ENABLED=false.'
       )
 
   is_cpu_device = device.platform == 'cpu'
-  if backend_impl == types.BackendImpl.C:
+  if impl == types.Impl.C:
     if not is_cpu_device:
       raise AssertionError(
-          f'C backend implementation requires a CPU device, got {device}.'
+          f'C implementation requires a CPU device, got {device}.'
       )
 
-  # NB: JAX backend works with any device.
+  # NB: JAX implementation works with any device.
 
 
-def _resolve_backend_impl_and_device(
-    backend_impl: Optional[Union[str, types.BackendImpl]],
+def _resolve_impl_and_device(
+    impl: Optional[Union[str, types.Impl]],
     device: Optional[jax.Device] = None,
-) -> Tuple[types.BackendImpl, jax.Device]:
-  """Resolves a backend implementation and device."""
-  if backend_impl:
-    backend_impl = types.BackendImpl(backend_impl)
+) -> Tuple[types.Impl, jax.Device]:
+  """Resolves a implementation and device."""
+  if impl:
+    impl = types.Impl(impl)
 
-  has_backend_impl, has_device = backend_impl is not None, device is not None
-  if (has_backend_impl, has_device) == (True, True):
+  has_impl, has_device = impl is not None, device is not None
+  if (has_impl, has_device) == (True, True):
     pass
-  elif (has_backend_impl, has_device) == (True, False):
-    device = _resolve_device(backend_impl)
-  elif (has_backend_impl, has_device) == (False, True):
-    backend_impl = _resolve_backend_impl(device)
+  elif (has_impl, has_device) == (True, False):
+    device = _resolve_device(impl)
+  elif (has_impl, has_device) == (False, True):
+    impl = _resolve_impl(device)
   else:
     device = jax.devices(jax.default_backend())[0]
     logging.info('Using JAX default device: %s.', device)
-    backend_impl = _resolve_backend_impl(device)
+    impl = _resolve_impl(device)
 
-  _check_backend_impl_device_compatibility(backend_impl, device)
-  return backend_impl, device  # pytype: disable=bad-return-type
+  _check_impl_device_compatibility(impl, device)
+  return impl, device  # pytype: disable=bad-return-type
 
 
 def _strip_weak_type(tree):
@@ -167,7 +167,7 @@ def _strip_weak_type(tree):
 
 def _put_option(
     o: mujoco.MjOption,
-    backend_impl: types.BackendImpl,
+    impl: types.Impl,
     impl_fields: Optional[dict[str, Any]] = None,
 ) -> types.Option:
   """Returns mjx.Option given mujoco.MjOption."""
@@ -195,7 +195,7 @@ def _put_option(
   fields['disableflags'] = types.DisableBit(o.disableflags)
   fields['enableflags'] = types.EnableBit(o.enableflags)
 
-  if backend_impl == types.BackendImpl.JAX:
+  if impl == types.Impl.JAX:
     has_fluid_params = o.density > 0 or o.viscosity > 0 or o.wind.any()
     implicitfast = o.integrator == mujoco.mjtIntegrator.mjINT_IMPLICITFAST
     if implicitfast and has_fluid_params:
@@ -203,12 +203,12 @@ def _put_option(
     fields['has_fluid_params'] = has_fluid_params
     return types.OptionJAX(**fields, **(impl_fields or {}))
 
-  if backend_impl == types.BackendImpl.C:
+  if impl == types.Impl.C:
     c_field_keys = types.OptionC.__annotations__.keys() - fields.keys()
     c_fields = {k: getattr(o, k, None) for k in c_field_keys}
     return types.OptionC(**fields, **c_fields, **(impl_fields or {}))
 
-  raise NotImplementedError(f'Unsupported backend: {backend_impl}')
+  raise NotImplementedError(f'Unsupported implementation: {impl}')
 
 
 def _put_statistic(s: mujoco.MjStatistic) -> types.Statistic:
@@ -283,7 +283,7 @@ def _put_model_jax(
   mj_field_names = {f.name for f in types.Model.fields() if f.name != '_impl'}
   fields = {f: getattr(m, f) for f in mj_field_names}
   fields['cam_mat0'] = fields['cam_mat0'].reshape((-1, 3, 3))
-  fields['opt'] = _put_option(m.opt, types.BackendImpl.JAX)
+  fields['opt'] = _put_option(m.opt, types.Impl.JAX)
   fields['stat'] = _put_statistic(m.stat)
 
   fields_jax = {}
@@ -340,7 +340,7 @@ def _put_model_c(
   mj_field_names = {f.name for f in types.Model.fields() if f.name != '_impl'}
   fields = {f: getattr(m, f) for f in mj_field_names}
   fields['cam_mat0'] = fields['cam_mat0'].reshape((-1, 3, 3))
-  fields['opt'] = _put_option(m.opt, backend_impl=types.BackendImpl.C)
+  fields['opt'] = _put_option(m.opt, impl=types.Impl.C)
   fields['stat'] = _put_statistic(m.stat)
 
   c_impl_keys = (
@@ -359,7 +359,7 @@ def _put_model_c(
 def put_model(
     m: mujoco.MjModel,
     device: Optional[jax.Device] = None,
-    backend_impl: Optional[Union[str, types.BackendImpl]] = None,
+    impl: Optional[Union[str, types.Impl]] = None,
     _full_compat: bool = False,  # pylint: disable=invalid-name
 ) -> types.Model:
   """Puts mujoco.MjModel onto a device, resulting in mjx.Model.
@@ -367,7 +367,7 @@ def put_model(
   Args:
     m: the model to put onto device
     device: which device to use - if unspecified picks the default device
-    backend_impl: backend implementation to use
+    impl: implementation to use
     _full_compat: put all MjModel fields onto device irrespective of MJX support
       This is an experimental feature.  Avoid using it for now.
 
@@ -375,28 +375,29 @@ def put_model(
     an mjx.Model placed on device
 
   Raises:
-    ValueError: if backend_impl is not supported
+    ValueError: if impl is not supported
     DeprecationWarning: if _full_compat is True
   """
 
   if _full_compat:
     warnings.warn(
-        'mjx.put_model(..., _full_compat=True) is deprecated.  Use'
-        ' mjx.put_model(..., backend_impl=types.BackendImpl.C) instead.',
+        'mjx.put_model(..., _full_compat=True) is deprecated and will be'
+        ' removed in MuJoCo >=3.4.  Use mjx.put_model(..., impl=types.Impl.C)'
+        ' instead.',
         DeprecationWarning,
         stacklevel=2,
     )
-    backend_impl = types.BackendImpl.C
+    impl = types.Impl.C
 
-  backend_impl, device = _resolve_backend_impl_and_device(backend_impl, device)
-  if backend_impl == types.BackendImpl.JAX:
+  impl, device = _resolve_impl_and_device(impl, device)
+  if impl == types.Impl.JAX:
     return _put_model_jax(m, device)
-  elif backend_impl == types.BackendImpl.C:
+  elif impl == types.Impl.C:
     return _put_model_c(m, device)
-  elif backend_impl == types.BackendImpl.WARP:
-    raise NotImplementedError('Warp backend not implemented yet.')
+  elif impl == types.Impl.WARP:
+    raise NotImplementedError('Warp implementation not implemented yet.')
   else:
-    raise ValueError(f'Unsupported backend implementation: {backend_impl}')
+    raise ValueError(f'Unsupported implementation: {impl}')
 
 
 def _make_data_public_fields(m: types.Model) -> Dict[str, Any]:
@@ -696,7 +697,7 @@ def _make_data_c(
 def make_data(
     m: Union[types.Model, mujoco.MjModel],
     device: Optional[jax.Device] = None,
-    backend_impl: Optional[Union[str, types.BackendImpl]] = None,
+    impl: Optional[Union[str, types.Impl]] = None,
     _full_compat: bool = False,  # pylint: disable=invalid-name
 ) -> types.Data:
   """Allocate and initialize Data.
@@ -704,7 +705,7 @@ def make_data(
   Args:
     m: the model to use
     device: which device to use - if unspecified picks the default device
-    backend_impl: backend implementation to use
+    impl: implementation to use ('jax', 'warp')
     _full_compat: put all fields onto device irrespective of MJX support This is
       an experimental feature.  Avoid using it for now. If using this flag, also
       use _full_compat for put_model.
@@ -713,35 +714,34 @@ def make_data(
     an initialized mjx.Data placed on device
 
   Raises:
-    ValueError: if the model's backend_impl does not match the make_data
-      backend_impl
-    NotImplementedError: if the backend_impl is not implemented yet
+    ValueError: if the model's impl does not match the make_data impl
+    NotImplementedError: if the impl is not implemented yet
     DeprecationWarning: if _full_compat is used
   """
   if _full_compat:
     warnings.warn(
         'mjx.make_data(..., _full_compat=True) is deprecated.  Use'
-        ' mjx.make_data(..., backend_impl=types.BackendImpl.C) instead.',
+        ' mjx.make_data(..., impl=types.Impl.C) instead.',
         DeprecationWarning,
         stacklevel=2,
     )
-    backend_impl = types.BackendImpl.C
+    impl = types.Impl.C
 
-  backend_impl, device = _resolve_backend_impl_and_device(backend_impl, device)
+  impl, device = _resolve_impl_and_device(impl, device)
 
-  if isinstance(m, types.Model) and m.backend_impl != backend_impl:
+  if isinstance(m, types.Model) and m.impl != impl:
     raise ValueError(
-        f'Model backend_impl {m.backend_impl} does not match make_data '
-        f'backend_impl {backend_impl}.'
+        f'Model impl {m.impl} does not match make_data '
+        f'implementation {impl}.'
     )
 
-  if backend_impl == types.BackendImpl.JAX:
+  if impl == types.Impl.JAX:
     return _make_data_jax(m, device)
-  elif backend_impl == types.BackendImpl.C:
+  elif impl == types.Impl.C:
     return _make_data_c(m, device)
 
   raise NotImplementedError(
-      f'make_data for backend_impl "{backend_impl}" not implemented yet.'
+      f'make_data for implementation "{impl}" not implemented yet.'
   )
 
 
@@ -951,7 +951,7 @@ def _put_data_c(
       if hasattr(d, f.name)
   }
 
-  # TODO(stunya): support islanding via C backend impl.
+  # TODO(stunya): support islanding via C impl.
   impl_fields['solver_niter'] = impl_fields['solver_niter'][0]
 
   # TODO(btaba): remove dense actuator moment.
@@ -1039,7 +1039,7 @@ def put_data(
     m: mujoco.MjModel,
     d: mujoco.MjData,
     device: Optional[jax.Device] = None,
-    backend_impl: Optional[Union[str, types.BackendImpl]] = None,
+    impl: Optional[Union[str, types.Impl]] = None,
     _full_compat: bool = False,  # pylint: disable=invalid-name
 ) -> types.Data:
   """Puts mujoco.MjData onto a device, resulting in mjx.Data.
@@ -1048,7 +1048,7 @@ def put_data(
     m: the model to use
     d: the data to put on device
     device: which device to use - if unspecified picks the default device
-    backend_impl: backend implementation to use
+    impl: implementation to use ('jax', 'warp')
     _full_compat: put all MjModel fields onto device irrespective of MJX support
       This is an experimental feature.  Avoid using it for now. If using this
       flag, also use _full_compat for put_model.
@@ -1059,20 +1059,20 @@ def put_data(
   if _full_compat:
     warnings.warn(
         'mjx.put_data(..., _full_compat=True) is deprecated.  Use'
-        ' mjx.put_data(..., backend_impl=types.BackendImpl.C) instead.',
+        ' mjx.put_data(..., impl=types.Impl.C) instead.',
         DeprecationWarning,
         stacklevel=2,
     )
-    backend_impl = types.BackendImpl.C
+    impl = types.Impl.C
 
-  backend_impl, device = _resolve_backend_impl_and_device(backend_impl, device)
-  if backend_impl == types.BackendImpl.JAX:
+  impl, device = _resolve_impl_and_device(impl, device)
+  if impl == types.Impl.JAX:
     return _put_data_jax(m, d, device)
-  elif backend_impl == types.BackendImpl.C:
+  elif impl == types.Impl.C:
     return _put_data_c(m, d, device)
 
   raise NotImplementedError(
-      f'put_data for backend_impl "{backend_impl}" not implemented yet.'
+      f'put_data for implementation "{impl}" not implemented yet.'
   )
 
 
@@ -1097,7 +1097,7 @@ def _get_data_into(
   batch_size = d.qpos.shape[0] if batched else 1
 
   dof_i, dof_j = [], []
-  if d.backend_impl == types.BackendImpl.JAX:
+  if d.impl == types.Impl.JAX:
     for i in range(m.nv):
       j = i
       while j > -1:
@@ -1116,13 +1116,13 @@ def _get_data_into(
     if ncon != result_i.ncon or nefc != result_i.nefc or nj != result_i.nJ:
       mujoco._functions._realloc_con_efc(result_i, ncon=ncon, nefc=nefc, nJ=nj)  # pylint: disable=protected-access
 
-    if d.backend_impl == types.BackendImpl.JAX:
+    if d.impl == types.Impl.JAX:
       all_fields = types.Data.fields() + types.DataJAX.fields()
-    elif d.backend_impl == types.BackendImpl.C:
+    elif d.impl == types.Impl.C:
       all_fields = types.Data.fields() + types.DataC.fields()
     else:
       raise NotImplementedError(
-          f'get_data_into for backend_impl "{d.backend_impl}" not implemented'
+          f'get_data_into for implementation "{d.impl}" not implemented'
           ' yet.'
       )
 
@@ -1188,7 +1188,7 @@ def _get_data_into(
           value = value.reshape(-1)
       elif field.name.startswith('efc_'):
         value = value[efc_active]
-      if d.backend_impl == types.BackendImpl.JAX:
+      if d.impl == types.Impl.JAX:
         if field.name == 'qM' and not support.is_sparse(m):
           value = value[dof_i, dof_j]
         elif field.name == 'qLD' and not support.is_sparse(m):
@@ -1226,12 +1226,12 @@ def get_data_into(
 
   d = jax.device_get(d)
 
-  if d.backend_impl in (types.BackendImpl.JAX, types.BackendImpl.C):
+  if d.impl in (types.Impl.JAX, types.Impl.C):
     # TODO(stunya): Split out _get_data_into once codepaths diverge enough.
     return _get_data_into(result, m, d)
 
   raise NotImplementedError(
-      f'get_data_into for backend_impl "{d.backend_impl}" not implemented yet.'
+      f'get_data_into for implementation "{d.impl}" not implemented yet.'
   )
 
 

mjx/mujoco/mjx/_src/io_test.py

@@ -26,8 +26,8 @@ from mujoco.mjx._src import io as mjx_io
 from mujoco.mjx._src import test_util
 
 # pylint: disable=g-importing-member
-from mujoco.mjx._src.types import BackendImpl
 from mujoco.mjx._src.types import ConeType
+from mujoco.mjx._src.types import Impl
 # pylint: enable=g-importing-member
 import numpy as np
 
@@ -114,11 +114,11 @@ class ModelIOTest(parameterized.TestCase):
 
   @parameterized.product(
       xml=(_MULTIPLE_CONVEX_OBJECTS, _MULTIPLE_CONSTRAINTS),
-      backend_impl=('jax', 'c'),
+      impl=('jax', 'c'),
   )
-  def test_put_model(self, xml, backend_impl):
+  def test_put_model(self, xml, impl):
     m = mujoco.MjModel.from_xml_string(xml)
-    mx = mjx.put_model(m, backend_impl=backend_impl)
+    mx = mjx.put_model(m, impl=impl)
 
     def assert_not_weak_type(x):
       if isinstance(x, jax.Array):
@@ -140,10 +140,10 @@ class ModelIOTest(parameterized.TestCase):
     self.assertEqual(mx.nM, m.nM)
     self.assertAlmostEqual(mx.opt.timestep, m.opt.timestep)
 
-    if backend_impl == 'jax':
+    if impl == 'jax':
       # fields restricted to MuJoCo should not be populated
       self.assertFalse(hasattr(mx, 'bvh_aabb'))
-    elif backend_impl == 'c':
+    elif impl == 'c':
       # Options specific to C are populated.
       self.assertEqual(mx.opt.apirate, m.opt.apirate)
       # Fields private to C backend impl are populated.
@@ -177,7 +177,7 @@ class ModelIOTest(parameterized.TestCase):
         mujoco.MjModel.from_xml_string(
             '<mujoco><option viscosity="3.0"/><worldbody/></mujoco>'
         ),
-        backend_impl='jax',
+        impl='jax',
     )
     self.assertTrue(m.opt.has_fluid_params)
 
@@ -218,7 +218,7 @@ class ModelIOTest(parameterized.TestCase):
             </body>
           </worldbody>
         </mujoco>"""),
-          backend_impl='jax',
+          impl='jax',
       )
 
   def test_implicitfast_fluid_not_implemented(self):
@@ -229,18 +229,18 @@ class ModelIOTest(parameterized.TestCase):
           <option viscosity="3.0" integrator="implicitfast"/>
           <worldbody/>
         </mujoco>"""),
-          backend_impl='jax',
+          impl='jax',
       )
 
   def test_wrap_inside(self):
     m = test_util.load_test_file('tendon/wrap_sidesite.xml')
-    mx0 = mjx.put_model(m, backend_impl='jax')
+    mx0 = mjx.put_model(m, impl='jax')
     np.testing.assert_equal(
         mx0._impl.is_wrap_inside,
         np.array([1, 0, 1, 0, 1, 1, 0]),
     )
     m.site_pos[2] = m.site_pos[1]
-    mx1 = mjx.put_model(m, backend_impl='jax')
+    mx1 = mjx.put_model(m, impl='jax')
     np.testing.assert_equal(
         mx1._impl.is_wrap_inside,
         np.array([0, 0, 1, 0, 1, 0, 0]),
@@ -251,10 +251,10 @@ class DataIOTest(parameterized.TestCase):
   """IO tests for mjx.Data."""
 
   @parameterized.parameters('jax', 'c')
-  def test_make_data(self, backend_impl: str):
+  def test_make_data(self, impl: str):
     """Test that make_data returns the correct shapes."""
     m = mujoco.MjModel.from_xml_string(_MULTIPLE_CONVEX_OBJECTS)
-    d = mjx.make_data(m, backend_impl=backend_impl)
+    d = mjx.make_data(m, impl=impl)
 
     nq = 22
     nbody = 5
@@ -312,34 +312,34 @@ class DataIOTest(parameterized.TestCase):
     self.assertEqual(d.qfrc_inverse.shape, (nv,))
     self.assertEqual(d._impl.efc_force.shape, (nefc,))
 
-    if backend_impl == 'jax':
+    if impl == 'jax':
       self.assertEqual(d._impl.qM.shape, (nv, nv))
       self.assertEqual(d._impl.qLD.shape, (nv, nv))
       self.assertEqual(d._impl.qLDiagInv.shape, (0,))
-    elif backend_impl == 'c':
+    elif impl == 'c':
       self.assertEqual(d._impl.qM.shape, (nm,))
       self.assertEqual(d._impl.qLD.shape, (nm,))
       self.assertEqual(d._impl.qLDiagInv.shape, (nv,))
 
     # test sparse
     m.opt.jacobian = mujoco.mjtJacobian.mjJAC_SPARSE
-    d = mjx.make_data(m, backend_impl=backend_impl)
+    d = mjx.make_data(m, impl=impl)
     self.assertEqual(d._impl.qM.shape, (nm,))
     self.assertEqual(d._impl.qLD.shape, (nm,))
     self.assertEqual(d._impl.qLDiagInv.shape, (nv,))
 
-    if backend_impl == 'c':
+    if impl == 'c':
       # check C specific fields
       self.assertEqual(d._impl.light_xpos.shape, (m.nlight, 3))
       self.assertEqual(d._impl.bvh_active.shape, (m.nbvh,))
 
   @parameterized.parameters('jax', 'c')
-  def test_put_data(self, backend_impl: str):
+  def test_put_data(self, impl: str):
     """Test that put_data puts the correct data for dense and sparse."""
     m = mujoco.MjModel.from_xml_string(_MULTIPLE_CONSTRAINTS)
     d = mujoco.MjData(m)
     mujoco.mj_step(m, d, 2)
-    dx = mjx.put_data(m, d, backend_impl=backend_impl)
+    dx = mjx.put_data(m, d, impl=impl)
 
     # check a few fields
     np.testing.assert_allclose(dx.qpos, d.qpos)
@@ -356,12 +356,12 @@ class DataIOTest(parameterized.TestCase):
         )
     )
 
-    if backend_impl == 'jax':
+    if impl == 'jax':
       # check that qM is transformed properly
       qm = np.zeros((m.nv, m.nv), dtype=np.float64)
       mujoco.mj_fullM(m, qm, d.qM)
       np.testing.assert_allclose(qm, mjx.full_m(mjx.put_model(m), dx))
-    elif backend_impl == 'c':
+    elif impl == 'c':
       np.testing.assert_allclose(dx._impl.qM, d.qM)
       np.testing.assert_allclose(dx._impl.qLD, d.qLD)
       np.testing.assert_allclose(dx._impl.qLDiagInv, d.qLDiagInv)
@@ -417,7 +417,7 @@ class DataIOTest(parameterized.TestCase):
     m.opt.jacobian = mujoco.mjtJacobian.mjJAC_SPARSE
     d = mujoco.MjData(m)
     mujoco.mj_step(m, d, 2)
-    dx_sparse = mjx.put_data(m, d, backend_impl=backend_impl)
+    dx_sparse = mjx.put_data(m, d, impl=impl)
     np.testing.assert_allclose(dx_sparse._impl.efc_J, dx._impl.efc_J, atol=1e-8)
 
     # check sparse mass matrices are correct
@@ -431,25 +431,25 @@ class DataIOTest(parameterized.TestCase):
     m.opt.jacobian = mujoco.mjtJacobian.mjJAC_DENSE
     d = mujoco.MjData(m)
     mujoco.mj_step(m, d, 2)
-    dx_from_dense = mjx.put_data(m, d, backend_impl=backend_impl)
-    if backend_impl == 'jax':
+    dx_from_dense = mjx.put_data(m, d, impl=impl)
+    if impl == 'jax':
       qm = np.zeros((m.nv, m.nv))
       mujoco.mj_fullM(m, qm, d.qM)
       np.testing.assert_allclose(dx_from_dense._impl.qM, qm, atol=1e-8)
-    elif backend_impl == 'c':
+    elif impl == 'c':
       np.testing.assert_allclose(dx_from_dense._impl.qM, d.qM, atol=1e-8)
 
   @parameterized.parameters(
       ('jax', False), ('jax', True), ('c', False), ('c', True)
   )
-  def test_get_data(self, backend_impl: str, sparse: bool):
+  def test_get_data(self, impl: str, sparse: bool):
     """Test that get_data makes correct MjData."""
     m = mujoco.MjModel.from_xml_string(_MULTIPLE_CONSTRAINTS)
     if sparse:
       m.opt.jacobian = mujoco.mjtJacobian.mjJAC_SPARSE
     d = mujoco.MjData(m)
     mujoco.mj_step(m, d, 2)
-    dx = mjx.put_data(m, d, backend_impl=backend_impl)
+    dx = mjx.put_data(m, d, impl=impl)
     d_2: mujoco.MjData = mjx.get_data(m, dx)
 
     # check a few fields
@@ -505,7 +505,7 @@ class DataIOTest(parameterized.TestCase):
     np.testing.assert_allclose(d_2.efc_aref, d.efc_aref)
     np.testing.assert_allclose(d_2.contact.efc_address, d.contact.efc_address)
 
-    if backend_impl == 'c':
+    if impl == 'c':
       # check fields specific to the C implementation
       np.testing.assert_allclose(d_2.bvh_active, d.bvh_active)
 
@@ -540,13 +540,13 @@ class DataIOTest(parameterized.TestCase):
     mjx.get_data(m, dx)
 
   @parameterized.parameters('jax', 'c')
-  def test_get_data_batched(self, backend_impl):
+  def test_get_data_batched(self, impl):
     """Test that get_data makes correct List[MjData] for batched Data."""
 
     m = mujoco.MjModel.from_xml_string(_MULTIPLE_CONSTRAINTS)
     d = mujoco.MjData(m)
     mujoco.mj_step(m, d, 2)
-    dx = mjx.put_data(m, d, backend_impl=backend_impl)
+    dx = mjx.put_data(m, d, impl=impl)
     # second data in batch has contact dist > 0, disables contact
     dx_b = jax.tree_util.tree_map(lambda x: jp.stack((x, x + 0.05)), dx)
     ds = mjx.get_data(m, dx_b)
@@ -557,13 +557,13 @@ class DataIOTest(parameterized.TestCase):
     self.assertEqual(ds[1].ncon, 0)
 
   @parameterized.parameters('jax', 'c')
-  def test_get_data_into(self, backend_impl):
+  def test_get_data_into(self, impl):
     """Test that get_data_into correctly populates an MjData."""
 
     m = mujoco.MjModel.from_xml_string(_MULTIPLE_CONSTRAINTS)
     d = mujoco.MjData(m)
     mujoco.mj_step(m, d, 2)
-    dx = mjx.put_data(m, d, backend_impl=backend_impl)
+    dx = mjx.put_data(m, d, impl=impl)
     d_2 = mujoco.MjData(m)
     mjx.get_data_into(d_2, m, dx)
 
@@ -580,32 +580,32 @@ class DataIOTest(parameterized.TestCase):
     np.testing.assert_allclose(d_2.contact.frame, d.contact.frame)
 
   @parameterized.parameters('jax', 'c')
-  def test_get_data_into_wrong_shape(self, backend_impl):
+  def test_get_data_into_wrong_shape(self, impl):
     """Tests that get_data_into throwsif input and output shapes don't match."""
 
     m = mujoco.MjModel.from_xml_string(_MULTIPLE_CONSTRAINTS)
     d = mujoco.MjData(m)
     mujoco.mj_step(m, d, 2)
-    dx = mjx.put_data(m, d, backend_impl=backend_impl)
+    dx = mjx.put_data(m, d, impl=impl)
     m_2 = mujoco.MjModel.from_xml_string(_MULTIPLE_CONVEX_OBJECTS)
     d_2 = mujoco.MjData(m_2)
     with self.assertRaisesRegex(ValueError, r'Input field.*has shape.*'):
       mjx.get_data_into(d_2, m, dx)
 
   @parameterized.parameters('jax', 'c')
-  def test_make_matches_put(self, backend_impl):
+  def test_make_matches_put(self, impl):
     """Test that make_data produces a pytree that matches put_data."""
     m = mujoco.MjModel.from_xml_string(_MULTIPLE_CONSTRAINTS)
     d = mujoco.MjData(m)
     mujoco.mj_step(m, d, 2)
-    dx = mjx.put_data(m, d, backend_impl=backend_impl)
+    dx = mjx.put_data(m, d, impl=impl)
 
     step_fn = lambda d: d.replace(time=d.time + 1)
     step_fn_jit = jax.jit(step_fn).lower(dx).compile()
 
     # placing an MjData onto device should yield the same treedef mjx.Data as
     # calling make_data.  they should be interchangeable for jax functions:
-    step_fn_jit(mjx.make_data(m, backend_impl=backend_impl))
+    step_fn_jit(mjx.make_data(m, impl=impl))
 
   def test_contact_elliptic_condim1(self):
     """Test that condim=1 with ConeType.ELLIPTIC is not implemented."""
@@ -654,7 +654,7 @@ class DataIOTest(parameterized.TestCase):
         </mujoco>
       """)
     with self.assertRaises(NotImplementedError):
-      mjx.put_model(m, backend_impl='jax')
+      mjx.put_model(m, impl='jax')
 
 
 class FullCompatTest(parameterized.TestCase):
@@ -678,71 +678,71 @@ class FullCompatTest(parameterized.TestCase):
     m = mujoco.MjModel.from_xml_string(xml)
     with self.assertWarns(DeprecationWarning):
       out = mjx_io.put_model(m, _full_compat=True)
-      self.assertEqual(out.backend_impl, BackendImpl.C)
+      self.assertEqual(out.impl, Impl.C)
     with self.assertWarns(DeprecationWarning):
       out = mjx_io.make_data(m, _full_compat=True)
-      self.assertEqual(out.backend_impl, BackendImpl.C)
+      self.assertEqual(out.impl, Impl.C)
 
 
-# Test cases for `_resolve_backend_impl_and_device` where the device is
+# Test cases for `_resolve_impl_and_device` where the device is
 # specified by the user and the device is available.
 _DEVICE_TEST_CASES = [
     # Arguments use the following format:
-    # (device_type_str, backend_impl_str,
-    #  (expected_device, expected_backend_impl)))
+    # (device_type_str, impl_str,
+    #  (expected_device, expected_impl)))
     # No backend specified.
-    ('cpu', None, ('cpu', BackendImpl.C)),
-    ('gpu-notnvidia', None, ('gpu', BackendImpl.JAX)),
-    ('gpu-nvidia', None, ('gpu', BackendImpl.WARP)),
-    ('tpu', None, ('tpu', BackendImpl.JAX)),
+    ('cpu', None, ('cpu', Impl.C)),
+    ('gpu-notnvidia', None, ('gpu', Impl.JAX)),
+    ('gpu-nvidia', None, ('gpu', Impl.WARP)),
+    ('tpu', None, ('tpu', Impl.JAX)),
     # JAX backend specified.
-    ('cpu', 'jax', ('cpu', BackendImpl.JAX)),
-    ('gpu-notnvidia', 'jax', ('gpu', BackendImpl.JAX)),
-    ('gpu-nvidia', 'jax', ('gpu', BackendImpl.JAX)),
-    ('tpu', 'jax', ('tpu', BackendImpl.JAX)),
+    ('cpu', 'jax', ('cpu', Impl.JAX)),
+    ('gpu-notnvidia', 'jax', ('gpu', Impl.JAX)),
+    ('gpu-nvidia', 'jax', ('gpu', Impl.JAX)),
+    ('tpu', 'jax', ('tpu', Impl.JAX)),
     # WARP backend specified.
     ('cpu', 'warp', ('cpu', 'error')),
     ('gpu-notnvidia', 'warp', ('cpu', 'error')),
-    ('gpu-nvidia', 'warp', ('gpu', BackendImpl.WARP)),
+    ('gpu-nvidia', 'warp', ('gpu', Impl.WARP)),
     ('tpu', 'warp', ('tpu', 'error')),
     # C backend specified.
-    ('cpu', 'c', ('cpu', BackendImpl.C)),
+    ('cpu', 'c', ('cpu', Impl.C)),
     ('gpu-notnvidia', 'c', ('cpu', 'error')),
     ('gpu-nvidia', 'c', ('cpu', 'error')),
     ('tpu', 'c', ('tpu', 'error')),
 ]
 
-# Test cases for `_resolve_backend_impl_and_device` where the user does NOT
+# Test cases for `_resolve_impl_and_device` where the user does NOT
 # specify a device. We mock the JAX default device.
 _DEFAULT_DEVICE_TEST_CASES = [
     # Arguments use the following format:
-    # (jax.default_device, backend_impl_str,
-    #  (expected_device, expected_backend_impl))
+    # (jax.default_device, impl_str,
+    #  (expected_device, expected_impl))
     # No backend impl specified.
-    ('cpu', None, ('cpu', BackendImpl.C)),
-    ('gpu-notnvidia', None, ('gpu', BackendImpl.JAX)),
-    ('gpu-nvidia', None, ('gpu', BackendImpl.WARP)),
-    ('tpu', None, ('tpu', BackendImpl.JAX)),
+    ('cpu', None, ('cpu', Impl.C)),
+    ('gpu-notnvidia', None, ('gpu', Impl.JAX)),
+    ('gpu-nvidia', None, ('gpu', Impl.WARP)),
+    ('tpu', None, ('tpu', Impl.JAX)),
     # JAX backend impl specified.
-    ('cpu', 'jax', ('cpu', BackendImpl.JAX)),
-    ('gpu-notnvidia', 'jax', ('gpu', BackendImpl.JAX)),
-    ('gpu-nvidia', 'jax', ('gpu', BackendImpl.JAX)),
-    ('tpu', 'jax', ('tpu', BackendImpl.JAX)),
+    ('cpu', 'jax', ('cpu', Impl.JAX)),
+    ('gpu-notnvidia', 'jax', ('gpu', Impl.JAX)),
+    ('gpu-nvidia', 'jax', ('gpu', Impl.JAX)),
+    ('tpu', 'jax', ('tpu', Impl.JAX)),
     # WARP backend impl specified.
     ('cpu', 'warp', ('cpu', 'error')),
     ('gpu-notnvidia', 'warp', ('cpu', 'error')),
-    ('gpu-nvidia', 'warp', ('gpu', BackendImpl.WARP)),
+    ('gpu-nvidia', 'warp', ('gpu', Impl.WARP)),
     ('tpu', 'warp', ('tpu', 'error')),
     # C backend impl specified, CPU should always be available.
-    ('cpu', 'c', ('cpu', BackendImpl.C)),
-    ('gpu-notnvidia', 'c', ('cpu', BackendImpl.C)),
-    ('gpu-nvidia', 'c', ('cpu', BackendImpl.C)),
-    ('tpu', 'c', ('cpu', BackendImpl.C)),
+    ('cpu', 'c', ('cpu', Impl.C)),
+    ('gpu-notnvidia', 'c', ('cpu', Impl.C)),
+    ('gpu-nvidia', 'c', ('cpu', Impl.C)),
+    ('tpu', 'c', ('cpu', Impl.C)),
 ]
 
 
-class ResolveBackendImplAndDeviceTest(parameterized.TestCase):
-  """Tests for the _resolve_backend_impl_and_device function."""
+class ResolveImplAndDeviceTest(parameterized.TestCase):
+  """Tests for the _resolve_impl_and_device function."""
 
   def setUp(self):
     super().setUp()
@@ -783,7 +783,7 @@ class ResolveBackendImplAndDeviceTest(parameterized.TestCase):
   def test_resolve_with_device(
       self,
       device_type_str,
-      backend_impl_str,
+      impl_str,
       expected,
   ):
     """Tests various combinations of device and backend impls."""
@@ -811,21 +811,21 @@ class ResolveBackendImplAndDeviceTest(parameterized.TestCase):
 
     self.mock_jax_devices.side_effect = devices_side_effect
 
-    expected_device, expected_backend_impl = expected
-    if expected_backend_impl == 'error':
+    expected_device, expected_impl = expected
+    if expected_impl == 'error':
       with self.assertRaises(AssertionError):
-        mjx_io._resolve_backend_impl_and_device(
-            backend_impl=backend_impl_str, device=input_device
+        mjx_io._resolve_impl_and_device(
+            impl=impl_str, device=input_device
         )
       return
 
-    actual_backend_impl, actual_device = (
-        mjx_io._resolve_backend_impl_and_device(
-            backend_impl=backend_impl_str, device=input_device
+    actual_impl, actual_device = (
+        mjx_io._resolve_impl_and_device(
+            impl=impl_str, device=input_device
         )
     )
 
-    self.assertEqual(actual_backend_impl, expected_backend_impl)
+    self.assertEqual(actual_impl, expected_impl)
     self.assertIsNotNone(actual_device)
     self.assertEqual(actual_device.platform, expected_device)
 
@@ -839,7 +839,7 @@ class ResolveBackendImplAndDeviceTest(parameterized.TestCase):
   def test_resolve_without_device(
       self,
       default_device_str,
-      backend_impl_str,
+      impl_str,
       expected,
   ):
     """Tests various combinations of jax.default_device and backend impls."""
@@ -878,32 +878,32 @@ class ResolveBackendImplAndDeviceTest(parameterized.TestCase):
         lambda: default_device_side_effect_str
     )
 
-    expected_device, expected_backend_impl = expected
+    expected_device, expected_impl = expected
     if (
-        expected_backend_impl == 'error'
+        expected_impl == 'error'
         and default_device_str != 'gpu-nvidia'
-        and backend_impl_str == 'warp'
+        and impl_str == 'warp'
     ):
       with self.assertRaisesRegex(RuntimeError, 'cuda backend not supported'):
-        mjx_io._resolve_backend_impl_and_device(
-            backend_impl=backend_impl_str, device=None
+        mjx_io._resolve_impl_and_device(
+            impl=impl_str, device=None
         )
       return
 
-    if expected_backend_impl == 'error':
+    if expected_impl == 'error':
       with self.assertRaises(AssertionError):
-        mjx_io._resolve_backend_impl_and_device(
-            backend_impl=backend_impl_str, device=None
+        mjx_io._resolve_impl_and_device(
+            impl=impl_str, device=None
         )
       return
 
-    actual_backend_impl, actual_device = (
-        mjx_io._resolve_backend_impl_and_device(
-            backend_impl=backend_impl_str, device=None
+    actual_impl, actual_device = (
+        mjx_io._resolve_impl_and_device(
+            impl=impl_str, device=None
         )
     )
 
-    self.assertEqual(actual_backend_impl, expected_backend_impl)
+    self.assertEqual(actual_impl, expected_impl)
     self.assertIsNotNone(actual_device)
     self.assertEqual(actual_device.platform, expected_device)
 
@@ -918,26 +918,26 @@ class ResolveBackendImplAndDeviceTest(parameterized.TestCase):
     self.mock_default_backend.side_effect = lambda: 'gpu'
 
     # Default to JAX instead of WARP on NVIDIA GPU.
-    backend_impl, device = mjx_io._resolve_backend_impl_and_device(
-        backend_impl=None, device=None
+    impl, device = mjx_io._resolve_impl_and_device(
+        impl=None, device=None
     )
-    self.assertEqual(backend_impl, BackendImpl.JAX)
+    self.assertEqual(impl, Impl.JAX)
     self.assertEqual(device.platform, 'gpu')
 
     # Specifying an NVIDIA GPU should still choose JAX.
-    backend_impl, device = mjx_io._resolve_backend_impl_and_device(
-        backend_impl=None, device=self.mock_nvidia_gpu
+    impl, device = mjx_io._resolve_impl_and_device(
+        impl=None, device=self.mock_nvidia_gpu
     )
-    self.assertEqual(backend_impl, BackendImpl.JAX)
+    self.assertEqual(impl, Impl.JAX)
     self.assertEqual(device.platform, 'gpu')
 
     # Requesting warp explicitly should fail since it is disabled.
     with self.assertRaises(AssertionError):
-      mjx_io._resolve_backend_impl_and_device(
-          backend_impl='warp', device=self.mock_nvidia_gpu
+      mjx_io._resolve_impl_and_device(
+          impl='warp', device=self.mock_nvidia_gpu
       )
     with self.assertRaises(AssertionError):
-      mjx_io._resolve_backend_impl_and_device(backend_impl='warp', device=None)
+      mjx_io._resolve_impl_and_device(impl='warp', device=None)
 
   @mock.patch.dict(os.environ, {'MJX_C_DEFAULT_ENABLED': 'false'})
   def test_resolve_c_disabled(self):
@@ -950,30 +950,30 @@ class ResolveBackendImplAndDeviceTest(parameterized.TestCase):
     self.mock_default_backend.side_effect = lambda: 'cpu'
 
     # Default to JAX instead of C on CPU.
-    backend_impl, device = mjx_io._resolve_backend_impl_and_device(
-        backend_impl=None, device=None
+    impl, device = mjx_io._resolve_impl_and_device(
+        impl=None, device=None
     )
-    self.assertEqual(backend_impl, BackendImpl.JAX)
+    self.assertEqual(impl, Impl.JAX)
     self.assertEqual(device.platform, 'cpu')
 
     # Specifing CPU should still choose JAX.
-    backend_impl, device = mjx_io._resolve_backend_impl_and_device(
-        backend_impl=None, device=self.mock_cpu
+    impl, device = mjx_io._resolve_impl_and_device(
+        impl=None, device=self.mock_cpu
     )
-    self.assertEqual(backend_impl, BackendImpl.JAX)
+    self.assertEqual(impl, Impl.JAX)
     self.assertEqual(device.platform, 'cpu')
 
     # Specifying C should choose C!
-    backend_impl, device = mjx_io._resolve_backend_impl_and_device(
-        backend_impl='c', device=None
+    impl, device = mjx_io._resolve_impl_and_device(
+        impl='c', device=None
     )
-    self.assertEqual(backend_impl, BackendImpl.C)
+    self.assertEqual(impl, Impl.C)
     self.assertEqual(device.platform, 'cpu')
 
-    backend_impl, device = mjx_io._resolve_backend_impl_and_device(
-        backend_impl='c', device=self.mock_cpu
+    impl, device = mjx_io._resolve_impl_and_device(
+        impl='c', device=self.mock_cpu
     )
-    self.assertEqual(backend_impl, BackendImpl.C)
+    self.assertEqual(impl, Impl.C)
     self.assertEqual(device.platform, 'cpu')
 
 

mjx/mujoco/mjx/_src/sensor.py

@@ -22,7 +22,7 @@ from mujoco.mjx._src import math
 from mujoco.mjx._src import ray
 from mujoco.mjx._src import smooth
 from mujoco.mjx._src import support
-from mujoco.mjx._src.types import BackendImpl
+from mujoco.mjx._src.types import Impl
 from mujoco.mjx._src.types import Data
 from mujoco.mjx._src.types import DataJAX
 from mujoco.mjx._src.types import DisableBit

mjx/mujoco/mjx/_src/types.py

@@ -24,8 +24,8 @@ from mujoco.mjx._src.dataclasses import PyTreeNode  # pylint: disable=g-importin
 import numpy as np
 
 
-class BackendImpl(enum.Enum):
-  """Backend implementation to use."""
+class Impl(enum.Enum):
+  """Implementation to use."""
 
   C = 'c'
   JAX = 'jax'
@@ -34,7 +34,7 @@ class BackendImpl(enum.Enum):
   @classmethod
   def _missing_(cls, value):
     # This method is called only when lookup by value fails
-    # (e.g., BackendImpl('JAX') fails initially because 'JAX' != 'jax')
+    # (e.g., Impl('JAX') fails initially because 'JAX' != 'jax')
     if not isinstance(value, str):
       return None
     for member in cls:
@@ -871,10 +871,10 @@ class Model(PyTreeNode):
   _impl: Union[ModelC, ModelJAX]
 
   @property
-  def backend_impl(self) -> BackendImpl:
+  def impl(self) -> Impl:
     return {
-        ModelC: BackendImpl.C,
-        ModelJAX: BackendImpl.JAX,
+        ModelC: Impl.C,
+        ModelJAX: Impl.JAX,
     }[type(self._impl)]
 
   def __getattr__(self, name: str):
@@ -1119,10 +1119,10 @@ class Data(PyTreeNode):
   _impl: Union[DataC, DataJAX]
 
   @property
-  def backend_impl(self) -> BackendImpl:
+  def impl(self) -> Impl:
     return {
-        DataC: BackendImpl.C,
-        DataJAX: BackendImpl.JAX,
+        DataC: Impl.C,
+        DataJAX: Impl.JAX,
     }[type(self._impl)]
 
   def __getattr__(self, name: str):