Convert qLD to CSR format.
PiperOrigin-RevId: 723955038 Change-Id: I30c3dc7f59739e89ae5fff8841432bc74717ec1b
This commit is contained in:
committed by
Copybara-Service
parent
cb1696eb34
commit
c27d3758c2
@@ -368,7 +368,7 @@ def make_data(
|
||||
'efc_aref': (nefc, float),
|
||||
'efc_force': (nefc, float),
|
||||
'_qM_sparse': (m.nM, float),
|
||||
'_qLD_sparse': (m.nM, float),
|
||||
'_qLD_sparse': (m.nC, float),
|
||||
'_qLDiagInv_sparse': (m.nv, float),
|
||||
}
|
||||
|
||||
|
||||
@@ -91,7 +91,10 @@ class SmoothTest(absltest.TestCase):
|
||||
_assert_eq(dx._qM_sparse, np.zeros(0), '_qM_sparse')
|
||||
# factor_m
|
||||
dx = jax.jit(mjx.factor_m)(mx, mjx.put_data(m, d))
|
||||
_assert_attr_eq(d, dx, 'qLD')
|
||||
qLDLegacy = np.zeros(mx.nM) # pylint:disable=invalid-name
|
||||
for i in range(m.nC):
|
||||
qLDLegacy[d.mapM2C[i]] = d.qLD[i]
|
||||
_assert_eq(qLDLegacy, dx.qLD, 'qLD')
|
||||
_assert_attr_eq(d, dx, 'qLDiagInv')
|
||||
_assert_eq(dx._qLD_sparse, np.zeros(0), '_qLD_sparse')
|
||||
_assert_eq(dx._qLDiagInv_sparse, np.zeros(0), '_qLDiagInv_sparse')
|
||||
|
||||
@@ -1348,7 +1348,7 @@ class Data(PyTreeNode):
|
||||
efc_aref: reference pseudo-acceleration (nefc,)
|
||||
efc_force: constraint force in constraint space (nefc,)
|
||||
_qM_sparse: qM in sparse representation (nM,)
|
||||
_qLD_sparse: qLD in sparse representation (nM,)
|
||||
_qLD_sparse: qLD in sparse representation (nC,)
|
||||
_qLDiagInv_sparse: qLDiagInv in sparse representation (nv,)
|
||||
""" # fmt: skip
|
||||
# constant sizes:
|
||||
|
||||
Reference in New Issue
Block a user