Check MJX with JAX backend for flex.

PiperOrigin-RevId: 782964265
Change-Id: I744d114569738ad0d5c4ee11ffba90bd50ef25ed

doc/mjx.rst

@@ -293,6 +293,8 @@ The following features are **unsupported**:
      - ``PGS``
    * - :ref:`Sensors <mjtSensor>`
      - ``PLUGIN``, ``USER``
+   * - Flex
+     - All
 
 .. _MjxSharpBits:
 

mjx/mujoco/mjx/_src/io.py

@@ -227,6 +227,9 @@ def _put_model_jax(
     device: Optional[jax.Device] = None,
 ) -> types.Model:
   """Puts mujoco.MjModel onto a device, resulting in mjx.Model."""
+  if m.nflex:
+    raise NotImplementedError('Flex not implemented for JAX backend.')
+
   mesh_geomid = set()
   for g1, g2, ip in collision_driver.geom_pairs(m):
     t1, t2 = m.geom_type[[g1, g2]]

mjx/mujoco/mjx/_src/io_test.py

@@ -967,6 +967,19 @@ class ResolveImplAndDeviceTest(parameterized.TestCase):
     self.assertEqual(impl, Impl.C)
     self.assertEqual(device.platform, 'cpu')
 
+  def test_flex_jax(self):
+    with self.assertRaises(NotImplementedError):
+      m = mujoco.MjModel.from_xml_string("""
+      <mujoco>
+        <worldbody>
+          <flexcomp name="flex" type="grid" dim="1" count="3 1 1" mass="1" spacing=".1 .1 .1">
+            <pin id="0"/>
+          </flexcomp>
+        </worldbody>
+      </mujoco>
+      """)
+      mjx.put_model(m, impl='jax')
+
 
 if __name__ == '__main__':
   absltest.main()