From 854112885354d342abb08f264de24c89cac836ea Mon Sep 17 00:00:00 2001
From: Taylor Howell <taylorhowell@google.com>
Date: Wed, 18 Feb 2026 04:03:22 -0800
Subject: [PATCH] Update _get_data_into_warp. Fixes #3109

PiperOrigin-RevId: 871790690
Change-Id: I225579f7559ca5c1de92af6606d02ff763f9f4d1
---
 mjx/mujoco/mjx/_src/io.py      |  3 ++-
 mjx/mujoco/mjx/_src/io_test.py | 16 ++++++++++++++++
 2 files changed, 18 insertions(+), 1 deletion(-)

diff --git a/mjx/mujoco/mjx/_src/io.py b/mjx/mujoco/mjx/_src/io.py
index 1e1101a5..e683633b 100644
--- a/mjx/mujoco/mjx/_src/io.py
+++ b/mjx/mujoco/mjx/_src/io.py
@@ -1486,10 +1486,11 @@ def _get_data_into_warp(
       if field.name in (
           'actuator_moment',
           'contact',
-          'efc_J',
           'qM',
           'qLD',
           'qLDiagInv',
+          'ten_J',
+          'flexedge_J',
       ):
         continue
       if field.name.startswith('efc_'):
diff --git a/mjx/mujoco/mjx/_src/io_test.py b/mjx/mujoco/mjx/_src/io_test.py
index b47ccb03..8382a6cf 100644
--- a/mjx/mujoco/mjx/_src/io_test.py
+++ b/mjx/mujoco/mjx/_src/io_test.py
@@ -744,6 +744,22 @@ class DataIOTest(parameterized.TestCase):
     self.assertEqual(d_2.contact.frame.shape, (1, 9))
     np.testing.assert_allclose(d_2.contact.frame, d.contact.frame)
 
+  def test_get_data_into_warp(self):
+    """Test get_data_into for impl='warp'."""
+
+    # TODO(taylorhowell): After put_data supports impl='warp' update test above
+    # and remove this test.
+    if not mjxw.WARP_INSTALLED:
+      self.skipTest('Warp is not installed.')
+    if not mjx_io.has_cuda_gpu_device():
+      self.skipTest('No CUDA GPU device.')
+
+    m = mujoco.MjModel.from_xml_string('<mujoco></mujoco>')
+    d = mujoco.MjData(m)
+    mx = mjx.put_model(m, impl='warp')
+    dx = mjx.make_data(m, impl='warp')
+    mjx.get_data_into(d, mx, dx)
+
   @parameterized.parameters('jax', 'c')
   def test_get_data_into_wrong_shape(self, impl):
     """Tests that get_data_into throwsif input and output shapes don't match."""