From 2a700c98aa2cfb5bb26c04a116d3e224ddc6de03 Mon Sep 17 00:00:00 2001
From: Taylor Howell <taylorhowell@google.com>
Date: Mon, 14 Jul 2025 10:48:17 -0700
Subject: [PATCH] Check MJX with JAX backend for flex.

PiperOrigin-RevId: 782964265
Change-Id: I744d114569738ad0d5c4ee11ffba90bd50ef25ed
---
 doc/mjx.rst                    |  2 ++
 mjx/mujoco/mjx/_src/io.py      |  3 +++
 mjx/mujoco/mjx/_src/io_test.py | 13 +++++++++++++
 3 files changed, 18 insertions(+)

diff --git a/doc/mjx.rst b/doc/mjx.rst
index 79b725d4..6bcd5dfe 100644
--- a/doc/mjx.rst
+++ b/doc/mjx.rst
@@ -293,6 +293,8 @@ The following features are **unsupported**:
      - ``PGS``
    * - :ref:`Sensors <mjtSensor>`
      - ``PLUGIN``, ``USER``
+   * - Flex
+     - All
 
 .. _MjxSharpBits:
 
diff --git a/mjx/mujoco/mjx/_src/io.py b/mjx/mujoco/mjx/_src/io.py
index e980afe5..f5a66995 100644
--- a/mjx/mujoco/mjx/_src/io.py
+++ b/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]]
diff --git a/mjx/mujoco/mjx/_src/io_test.py b/mjx/mujoco/mjx/_src/io_test.py
index 706d7e5e..e3ecd539 100644
--- a/mjx/mujoco/mjx/_src/io_test.py
+++ b/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()