From 5b245e5d02cc201b51cfcb29c54ae74945d5214c Mon Sep 17 00:00:00 2001
From: Yuval Tassa <tassa@google.com>
Date: Wed, 29 Jan 2025 15:26:18 -0800
Subject: [PATCH] Fix handling of _full_compat argument in `put_data`.

PiperOrigin-RevId: 721128021
Change-Id: Ic1cce42fbb38799228d6433ef0c761b1f9b81c4b
---
 mjx/mujoco/mjx/_src/io.py | 25 ++++++++++++-------------
 1 file changed, 12 insertions(+), 13 deletions(-)

diff --git a/mjx/mujoco/mjx/_src/io.py b/mjx/mujoco/mjx/_src/io.py
index 89ed3804..6acef205 100644
--- a/mjx/mujoco/mjx/_src/io.py
+++ b/mjx/mujoco/mjx/_src/io.py
@@ -123,7 +123,6 @@ def put_model(
       if t == mujoco.mjtGeom.mjGEOM_MESH:
         mesh_geomid.add(g)
 
-
   # check for unsupported sensor and equality constraint combinations
   sensor_rne_postconstraint = (
       np.any(m.sensor_type == types.SensorType.ACCELEROMETER)
@@ -659,18 +658,6 @@ def put_data(
 
     fields[fname] = value
 
-  # convert qM and qLD if jacobian is dense
-  if not support.is_sparse(m):
-    fields['qM'] = np.zeros((m.nv, m.nv))
-    mujoco.mj_fullM(m, fields['qM'], d.qM)
-    # TODO(erikfrey): derive L*L' from L'*D*L instead of recomputing
-    try:
-      fields['qLD'], _ = scipy.linalg.cho_factor(fields['qM'])
-    except scipy.linalg.LinAlgError:
-      # this happens when qM is empty or unstable simulation
-      fields['qLD'] = np.zeros((m.nv, m.nv))
-    fields['qLDiagInv'] = np.zeros(0)
-
   if _full_compat:
     # full compatibility mode, we store sparse qM regardless of jacobian setting
     fields['_qM_sparse'] = fields['qM']
@@ -687,6 +674,18 @@ def put_data(
       ):
         fields[f.name] = np.zeros(0, dtype=fields[f.name].dtype)
 
+  # convert qM and qLD if jacobian is dense
+  if not support.is_sparse(m):
+    fields['qM'] = np.zeros((m.nv, m.nv))
+    mujoco.mj_fullM(m, fields['qM'], d.qM)
+    # TODO(erikfrey): derive L*L' from L'*D*L instead of recomputing
+    try:
+      fields['qLD'], _ = scipy.linalg.cho_factor(fields['qM'])
+    except scipy.linalg.LinAlgError:
+      # this happens when qM is empty or unstable simulation
+      fields['qLD'] = np.zeros((m.nv, m.nv))
+    fields['qLDiagInv'] = np.zeros(0)
+
   fields['contact'] = contact
   fields.update(ne=ne, nf=nf, nl=nl, nefc=nefc, ncon=ncon, efc_type=efc_type)