Remove mjData.qLDiagSqrtInv, add corresponding argument to mj_solveM2.

- `qLDiagSqrtInv` is only required for the dual solvers. It is now computed as-needed rather than unconditionally.
- `mj_solveM2` now requires a new input array `sqrtInvD` which contains the square root of the inverse diagonal D (formerly saved in `qLDiagSqrtInv`).

PiperOrigin-RevId: 710805133
Change-Id: I0622d6a8da3882916824e9c10bad9223c122c321

mjx/mujoco/mjx/_src/io.py

@@ -310,7 +310,6 @@ def make_data(
         'qM': (m.nM, float) if support.is_sparse(m) else (m.nv, m.nv, float),
         'qLD': (m.nM, float) if support.is_sparse(m) else (m.nv, m.nv, float),
         'qLDiagInv': (m.nv, float) if support.is_sparse(m) else (0, float),
-        'qLDiagSqrtInv': (m.nv, float),
         'bvh_aabb_dyn': (m.nbvhdynamic, 6, float),
         'bvh_active': (m.nbvh, jp.uint8),
         'flexedge_velocity': (m.nflexedge, float),

mjx/mujoco/mjx/_src/types.py

@@ -1274,7 +1274,6 @@ class Data(PyTreeNode):
                                                      if dense:  (nv, nv)
     qLDiagInv: 1/diag(D)                             if sparse: (nv,)
                                                      if dense:  (0,)
-    qLDiagSqrtInv: 1/sqrt(diag(D))                              (nv,)
     bvh_aabb_dyn: global bounding box (center, size)            (nbvhdynamic, 6)
     bvh_active: volume has been added to collisions             (nbvh,)
     flexedge_velocity: flex edge velocities                     (nflexedge,)
@@ -1404,7 +1403,6 @@ class Data(PyTreeNode):
   qM: jax.Array  # pylint:disable=invalid-name
   qLD: jax.Array  # pylint:disable=invalid-name
   qLDiagInv: jax.Array  # pylint:disable=invalid-name
-  qLDiagSqrtInv: jax.Array  # pylint:disable=invalid-name
   bvh_aabb_dyn: jax.Array = _restricted_to('mujoco')
   bvh_active: jax.Array = _restricted_to('mujoco')
   # position, velocity dependent: