chenlin/Mujoco_WASM
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Import google-deepmind/mujoco_warp from GitHub.

Browse Source 
PiperOrigin-RevId: 929061622
Change-Id: I53892a819c3727b1f04db943839606ec1f477df1
This commit is contained in:
Taylor Howell
2026-06-09 02:27:41 -07:00
committed by Copybara-Service
parent 58a8c6ee56
commit 763e713d1b
8 changed files with 425 additions and 605 deletions
Download Patch File Download Diff File Expand all files Collapse all files
mjx/mujoco/mjx/_src/io_test.py
-2
View File
@@ -182,8 +182,6 @@ class ModelIOTest(parameterized.TestCase):
self.assertFalse(hasattr(mx, 'bvh_aabb'))
elif impl == 'warp':
# Options specific to Warp are populated.
self.assertTrue(hasattr(mx.opt._impl, 'ls_parallel'))
# Fields private to Warp backend impl are populated.
self.assertTrue(hasattr(mx._impl, 'nxn_geom_pair'))
elif impl == 'cpp':
mjx/mujoco/mjx/third_party/mujoco_warp/_src/io.py
Vendored
+31 -5
View File
@@ -257,9 +257,6 @@ def put_model(mjm: mujoco.MjModel) -> types.Model:
opt.tolerance = max(opt.tolerance, 1e-6)
# warp only fields
ls_parallel_id = mujoco.mj_name2id(mjm, mujoco.mjtObj.mjOBJ_NUMERIC, "ls_parallel")
opt.ls_parallel = (ls_parallel_id > -1) and (mjm.numeric_data[mjm.numeric_adr[ls_parallel_id]] == 1)
opt.ls_parallel_min_step = 1.0e-6 # TODO(team): determine good default setting
opt.broadphase = types.BroadphaseType.NXN
opt.broadphase_filter = types.BroadphaseFilter.PLANE | types.BroadphaseFilter.SPHERE | types.BroadphaseFilter.OBB
opt.graph_conditional = True
@@ -301,6 +298,12 @@ def put_model(mjm: mujoco.MjModel) -> types.Model:
m.nmaxpyramid = np.maximum(1, 2 * (m.nmaxcondim - 1))
m.has_sdf_geom = (mjm.geom_type == mujoco.mjtGeom.mjGEOM_SDF).any()
m.block_dim = types.BlockDim()
# Derive CG solver block_dim from nv: clamp(round_up_to_32(nv), 32, 256)
_nv_block = max(32, min(256, ((mjm.nv + 31) // 32) * 32))
m.block_dim.update_gradient_grad = _nv_block
m.block_dim.solve_beta_accumulate = _nv_block
m.block_dim.solve_search_update_cg = _nv_block
m.block_dim.solve_init_search_cg = _nv_block
if mjm.nv > 500:
m.block_dim.linesearch_iterative = 512
m.is_sparse = is_sparse(mjm)
@@ -308,6 +311,26 @@ def put_model(mjm: mujoco.MjModel) -> types.Model:
m.max_ten_J_rownnz = int(mjm.ten_J_rownnz.max()) if mjm.ntendon else 0
# Upper bound on a contact's Jacobian support, to size the elliptic-cone JTCJ launch (one
# thread per (contact, support-pair)). A contact's row spans the dof chains of weld(b1) and
# weld(b2) (see _efc_contact_jac_sparse in constraint.py); take the largest union over all
# geom-carrying bodies -- a safe superset, since over-estimating only adds skipped threads.
# A body's dof chain is exactly the sparsity of its deepest dof's row in the (ancestor-
# structured) mass matrix, so reuse MuJoCo's precomputed M_colind rather than re-walking.
def _dof_chain(body):
if mjm.body_dofnum[body] == 0:
return frozenset()
dof = int(mjm.body_dofadr[body] + mjm.body_dofnum[body] - 1)
adr = int(mjm.M_rowadr[dof])
return frozenset(int(mjm.M_colind[adr + k]) for k in range(int(mjm.M_rownnz[dof])))
chains = list({_dof_chain(int(mjm.body_weldid[b])) for b in mjm.geom_bodyid})
max_rownnz = 0
for i, chain_i in enumerate(chains):
for chain_j in chains[i:]:
max_rownnz = max(max_rownnz, len(chain_i | chain_j))
m.jtcj_max_pairs = max(max_rownnz * (max_rownnz + 1) // 2, 1)
# body ids grouped by tree level (depth-based traversal)
bodies, body_depth = {}, np.zeros(mjm.nbody, dtype=int) - 1
for i in range(mjm.nbody):
@@ -2841,7 +2864,6 @@ def override_model(model: types.Model | mujoco.MjModel, overrides: dict[str, Any
Overrides are of the format:
opt.iterations = 1
opt.ls_parallel = True
opt.cone = pyramidal
opt.disableflags = contact | spring
"""
@@ -2865,7 +2887,6 @@ def override_model(model: types.Model | mujoco.MjModel, overrides: dict[str, Any
mjw_only_fields = {
"opt.broadphase",
"opt.broadphase_filter",
"opt.ls_parallel",
"opt.graph_conditional",
"opt.contact_sensor_maxmatch",
}
@@ -2881,6 +2902,11 @@ def override_model(model: types.Model | mujoco.MjModel, overrides: dict[str, Any
overrides = overrides_dict
for key, val in overrides.items():
if key == "opt.ls_parallel":
raise ValueError("ls_parallel was removed in MuJoCo Warp 3.9.1.")
if key == "opt.ls_parallel_min_step":
raise ValueError("ls_parallel_min_step was removed in MuJoCo Warp 3.9.1.")
# skip overrides on MjModel for properties that are only on mjw.Model
if key in mjw_only_fields and isinstance(model, mujoco.MjModel):
continue
mjx/mujoco/mjx/third_party/mujoco_warp/_src/passive.py
Vendored
+14 -6
View File
@@ -630,6 +630,12 @@ def _flex_elasticity(
f = i
break
stiffness_adr_base = flex_stiffnessadr[f]
if stiffness_adr_base < 0:
return
if flex_stiffness[stiffness_adr_base] == 0.0:
return
local_elemid = elemid - flex_elemadr[f]
dim = flex_dim[f]
nvert = dim + 1
@@ -671,7 +677,7 @@ def _flex_elasticity(
metric = wp.matrix(0.0, shape=(6, 6))
stiffness_size = nedge * (nedge + 1) / 2
stiffness_adr = flex_stiffnessadr[f] + local_elemid * stiffness_size
stiffness_adr = stiffness_adr_base + local_elemid * stiffness_size
id = int(0)
for ed1 in range(nedge):
for ed2 in range(ed1, nedge):
@@ -721,6 +727,10 @@ def _flex_bending(
f = i
break
bendingadr = flex_bendingadr[f]
if bendingadr < 0:
return
if flex_dim[f] != 2:
return
@@ -734,10 +744,8 @@ def _flex_bending(
flex_vertadr[f] + flex_edgeflap[edgeid][1],
)
adr = flex_bendingadr[f]
frc = wp.matrix(0.0, shape=(4, 3))
if flex_bending[adr + 16]:
if flex_bending[bendingadr + 16]:
v0 = flexvert_xpos_in[worldid, v[0]]
v1 = flexvert_xpos_in[worldid, v[1]]
v2 = flexvert_xpos_in[worldid, v[2]]
@@ -752,8 +760,8 @@ def _flex_bending(
for x in range(3):
acc = float(0.0)
for j in range(nvert):
acc += flex_bending[adr + 4 * i + j] * flexvert_xpos_in[worldid, v[j]][x]
force[i, x] = -(acc + flex_bending[adr + 16] * frc[i, x])
acc += flex_bending[bendingadr + 4 * i + j] * flexvert_xpos_in[worldid, v[j]][x]
force[i, x] = -(acc + flex_bending[bendingadr + 16] * frc[i, x])
for i in range(nvert):
bodyid = flex_vertbodyid[v[i]]
mjx/mujoco/mjx/third_party/mujoco_warp/_src/solver.py
Vendored
+316 -535
View File
File diff suppressed because it is too large Load Diff
mjx/mujoco/mjx/third_party/mujoco_warp/_src/types.py
Vendored
+27 -5
View File
@@ -66,6 +66,10 @@ class BlockDim:
update_gradient_JTDAJ_sparse: update gradient JTDAJ sparse block dimension (solver)
update_gradient_JTDAJ_dense: update gradient JTDAJ dense block dimension (solver)
linesearch_iterative: linesearch iterative block dimension (solver)
update_gradient_grad: update gradient grad block dimension (solver)
solve_beta_accumulate: solve beta accumulate block dimension (solver)
solve_search_update_cg: solve search update CG block dimension (solver)
solve_init_search_cg: solve init search CG block dimension (solver)
contact_jac_tiled: contact Jacobian tiled block dimension (solver)
qderiv_actuator_dense: qderiv actuator dense block dimension (derivative)
render: render block dimension (render)
@@ -92,6 +96,10 @@ class BlockDim:
update_gradient_JTDAJ_sparse: int = 64
update_gradient_JTDAJ_dense: int = 128
linesearch_iterative: int = 32
update_gradient_grad: int = 256
solve_beta_accumulate: int = 256
solve_search_update_cg: int = 256
solve_init_search_cg: int = 256
contact_jac_tiled: int = 32
# derivative
qderiv_actuator_dense: int = 32
@@ -800,8 +808,6 @@ class Option:
warp only fields:
impratio_invsqrt: ratio of friction-to-normal contact impedance (stored as inverse square root)
ls_parallel: evaluate engine solver step sizes in parallel
ls_parallel_min_step: minimum step size for solver linesearch
broadphase: broadphase type (BroadphaseType)
broadphase_filter: broadphase filter bitflag (BroadphaseFilter)
graph_conditional: flag to use cuda graph conditional
@@ -834,14 +840,29 @@ class Option:
sdf_iterations: int
# warp only fields:
impratio_invsqrt: array("*", float)
ls_parallel: bool
ls_parallel_min_step: float
broadphase: BroadphaseType
broadphase_filter: BroadphaseFilter
graph_conditional: bool
run_collision_detection: bool
contact_sensor_maxmatch: int
# TODO(team): remove in future version
@property
def ls_parallel(self) -> bool:
raise AttributeError("ls_parallel was removed in MuJoCo Warp 3.9.1.")
@ls_parallel.setter
def ls_parallel(self, value: bool):
raise AttributeError("ls_parallel was removed in MuJoCo Warp 3.9.1.")
@property
def ls_parallel_min_step(self) -> float:
raise AttributeError("ls_parallel_min_step was removed in MuJoCo Warp 3.9.1.")
@ls_parallel_min_step.setter
def ls_parallel_min_step(self, value: float):
raise AttributeError("ls_parallel_min_step was removed in MuJoCo Warp 3.9.1.")
@dataclasses.dataclass
class Statistic:
@@ -1290,6 +1311,7 @@ class Model:
tendon_geom_adr: geom tendon address
tendon_limited_adr: addresses for limited tendons
max_ten_J_rownnz: maximum number of non-zeros in a tendon row
jtcj_max_pairs: bound on a contact's support-pair count, sizes the elliptic-cone JTCJ launch
ten_wrapadr_site: wrap object starting address for sites
ten_wrapnum_site: number of site wrap objects per tendon
wrap_jnt_adr: addresses for joint tendon wrap object
@@ -1719,6 +1741,7 @@ class Model:
tendon_geom_adr: wp.array[int]
tendon_limited_adr: wp.array[int]
max_ten_J_rownnz: int
jtcj_max_pairs: int
ten_wrapadr_site: wp.array[int]
ten_wrapnum_site: wp.array[int]
wrap_jnt_adr: wp.array[int]
@@ -2205,7 +2228,6 @@ class SolverContext:
mv: wp.array2d[float]
jv: wp.array2d[float]
quad: wp.array2d[wp.vec3]
quad_gauss: wp.array[wp.vec3]
alpha: wp.array[float]
improvement: wp.array[float]
prev_grad: wp.array2d[float]
mjx/mujoco/mjx/third_party/mujoco_warp/viewer.py
Vendored
+1 -1
View File
@@ -169,7 +169,7 @@ def _main(argv: Sequence[str]) -> None:
solver, cone = mjw.SolverType(m.opt.solver).name, mjw.ConeType(m.opt.cone).name
integrator = mjw.IntegratorType(m.opt.integrator).name
iterations, ls_iterations = m.opt.iterations, m.opt.ls_iterations
ls_str = f"{'parallel' if m.opt.ls_parallel else 'iterative'} linesearch iterations: {ls_iterations}"
ls_str = f"linesearch iterations: {ls_iterations}"
print(
f" nbody: {m.nbody} nv: {m.nv} ngeom: {m.ngeom} nu: {m.nu} is_sparse: {m.is_sparse}\n"
f" broadphase: {broadphase} broadphase_filter: {filter}\n"
mjx/mujoco/mjx/warp/forward.py
+8 -17
View File
@@ -14,17 +14,14 @@
# ==============================================================================
"""DO NOT EDIT. This file is auto-generated."""
import dataclasses
import functools
import jax
import warp as wp
from mujoco.mjx._src import types
from mujoco.mjx.warp import ffi
import mujoco.mjx.third_party.mujoco_warp as mjwarp
from mujoco.mjx.third_party.mujoco_warp._src import types as mjwp_types
from mujoco.mjx.warp import ffi
import warp as wp
_m = mjwarp.Model(
**{f.name: None for f in dataclasses.fields(mjwarp.Model) if f.init}
@@ -241,6 +238,7 @@ def _forward_shim(
jnt_stiffness: wp.array2d[float],
jnt_stiffnesspoly: wp.array2d[wp.vec2],
jnt_type: wp.array[int],
jtcj_max_pairs: int,
light_bodyid: wp.array[int],
light_dir: wp.array2d[wp.vec3],
light_dir0: wp.array2d[wp.vec3],
@@ -415,8 +413,6 @@ def _forward_shim(
opt__integrator: int,
opt__iterations: int,
opt__ls_iterations: int,
opt__ls_parallel: bool,
opt__ls_parallel_min_step: float,
opt__ls_tolerance: wp.array[float],
opt__magnetic: wp.array[wp.vec3],
opt__run_collision_detection: bool,
@@ -780,6 +776,7 @@ def _forward_shim(
_m.jnt_stiffness = jnt_stiffness
_m.jnt_stiffnesspoly = jnt_stiffnesspoly
_m.jnt_type = jnt_type
_m.jtcj_max_pairs = jtcj_max_pairs
_m.light_bodyid = light_bodyid
_m.light_dir = light_dir
_m.light_dir0 = light_dir0
@@ -865,8 +862,6 @@ def _forward_shim(
_m.opt.integrator = opt__integrator
_m.opt.iterations = opt__iterations
_m.opt.ls_iterations = opt__ls_iterations
_m.opt.ls_parallel = opt__ls_parallel
_m.opt.ls_parallel_min_step = opt__ls_parallel_min_step
_m.opt.ls_tolerance = opt__ls_tolerance
_m.opt.magnetic = opt__magnetic
_m.opt.run_collision_detection = opt__run_collision_detection
@@ -1796,6 +1791,7 @@ def _forward_jax_impl(m: types.Model, d: types.Data):
m.jnt_stiffness,
m.jnt_stiffnesspoly,
m.jnt_type,
m._impl.jtcj_max_pairs,
m._impl.light_bodyid,
m.light_dir,
m.light_dir0,
@@ -1970,8 +1966,6 @@ def _forward_jax_impl(m: types.Model, d: types.Data):
m.opt.integrator,
m.opt.iterations,
m.opt.ls_iterations,
m.opt._impl.ls_parallel,
m.opt._impl.ls_parallel_min_step,
m.opt.ls_tolerance,
m.opt.magnetic,
m.opt._impl.run_collision_detection,
@@ -2498,6 +2492,7 @@ def _step_shim(
jnt_stiffness: wp.array2d[float],
jnt_stiffnesspoly: wp.array2d[wp.vec2],
jnt_type: wp.array[int],
jtcj_max_pairs: int,
light_bodyid: wp.array[int],
light_dir: wp.array2d[wp.vec3],
light_dir0: wp.array2d[wp.vec3],
@@ -2677,8 +2672,6 @@ def _step_shim(
opt__integrator: int,
opt__iterations: int,
opt__ls_iterations: int,
opt__ls_parallel: bool,
opt__ls_parallel_min_step: float,
opt__ls_tolerance: wp.array[float],
opt__magnetic: wp.array[wp.vec3],
opt__run_collision_detection: bool,
@@ -3048,6 +3041,7 @@ def _step_shim(
_m.jnt_stiffness = jnt_stiffness
_m.jnt_stiffnesspoly = jnt_stiffnesspoly
_m.jnt_type = jnt_type
_m.jtcj_max_pairs = jtcj_max_pairs
_m.light_bodyid = light_bodyid
_m.light_dir = light_dir
_m.light_dir0 = light_dir0
@@ -3136,8 +3130,6 @@ def _step_shim(
_m.opt.integrator = opt__integrator
_m.opt.iterations = opt__iterations
_m.opt.ls_iterations = opt__ls_iterations
_m.opt.ls_parallel = opt__ls_parallel
_m.opt.ls_parallel_min_step = opt__ls_parallel_min_step
_m.opt.ls_tolerance = opt__ls_tolerance
_m.opt.magnetic = opt__magnetic
_m.opt.run_collision_detection = opt__run_collision_detection
@@ -4089,6 +4081,7 @@ def _step_jax_impl(m: types.Model, d: types.Data):
m.jnt_stiffness,
m.jnt_stiffnesspoly,
m.jnt_type,
m._impl.jtcj_max_pairs,
m._impl.light_bodyid,
m.light_dir,
m.light_dir0,
@@ -4268,8 +4261,6 @@ def _step_jax_impl(m: types.Model, d: types.Data):
m.opt.integrator,
m.opt.iterations,
m.opt.ls_iterations,
m.opt._impl.ls_parallel,
m.opt._impl.ls_parallel_min_step,
m.opt.ls_tolerance,
m.opt.magnetic,
m.opt._impl.run_collision_detection,
mjx/mujoco/mjx/warp/types.py
+28 -34
View File
@@ -15,17 +15,14 @@
"""MJX Warp types.
DO NOT EDIT. This file is auto-generated.
"""
import dataclasses
import typing
from typing import Tuple
import jax
from jax import tree_util
from jax.interpreters import batching
import numpy as np
from mujoco.mjx._src import dataclasses as mjx_dataclasses
import numpy as np
if typing.TYPE_CHECKING:
GraphMode = int
@@ -37,7 +34,6 @@ if typing.TYPE_CHECKING:
else:
try:
from warp._src.jax_experimental.ffi import GraphMode
from mujoco.mjx.third_party.mujoco_warp._src import types as mjwp_types
Callback = mjwp_types.Callback
@@ -46,7 +42,6 @@ else:
Callback = None
PyTreeNode = mjx_dataclasses.PyTreeNode
@dataclasses.dataclass(frozen=True)
@tree_util.register_pytree_node_class
class TileSet:
@@ -58,12 +53,9 @@ class TileSet:
adr: address of each tile in the set
size: size of all the tiles in this set
"""
adr: np.ndarray
size: int
# Manually kept in this generated shim until TileSet method generation is
# needed more broadly. Keep this in sync with mujoco_warp._src.types.TileSet.
def __eq__(self, other) -> bool:
if self.__class__ is not other.__class__:
return NotImplemented
@@ -101,18 +93,25 @@ class BlockDim:
energy_vel_kinetic: energy velocity kinetic block dimension (sensor)
cholesky_factorize: Cholesky factorize block dimension (smooth)
cholesky_solve: Cholesky solve block dimension (smooth)
cholesky_factorize_solve: Cholesky factorize and solve block dimension (smooth)
cholesky_factorize_solve: Cholesky factorize and solve block dimension
(smooth)
solve_LD_sparse_fused: solve LD sparse fused block dimension (smooth)
update_gradient_cholesky: update gradient Cholesky block dimension (solver)
update_gradient_cholesky_blocked: update gradient Cholesky blocked block dimension (solver)
update_gradient_JTDAJ_sparse: update gradient JTDAJ sparse block dimension (solver)
update_gradient_JTDAJ_dense: update gradient JTDAJ dense block dimension (solver)
update_gradient_cholesky_blocked: update gradient Cholesky blocked block
dimension (solver)
update_gradient_JTDAJ_sparse: update gradient JTDAJ sparse block dimension
(solver)
update_gradient_JTDAJ_dense: update gradient JTDAJ dense block dimension
(solver)
linesearch_iterative: linesearch iterative block dimension (solver)
update_gradient_grad: update gradient grad block dimension (solver)
solve_beta_accumulate: solve beta accumulate block dimension (solver)
solve_search_update_cg: solve search update CG block dimension (solver)
solve_init_search_cg: solve init search CG block dimension (solver)
contact_jac_tiled: contact Jacobian tiled block dimension (solver)
qderiv_actuator_dense: qderiv actuator dense block dimension (derivative)
render: render block dimension (render)
"""
actuator_velocity: int
cholesky_factorize: int
cholesky_factorize_solve: int
@@ -127,10 +126,14 @@ class BlockDim:
render: int
segmented_sort: int
solve_LD_sparse_fused: int
solve_beta_accumulate: int
solve_init_search_cg: int
solve_search_update_cg: int
update_gradient_JTDAJ_dense: int
update_gradient_JTDAJ_sparse: int
update_gradient_cholesky: int
update_gradient_cholesky_blocked: int
update_gradient_grad: int
def tree_flatten(self):
children = list((getattr(self, k) for k in self.__dataclass_fields__))
@@ -144,13 +147,10 @@ class BlockDim:
class StatisticWarp(PyTreeNode):
"""Derived fields from Statistic."""
meaninertia: jax.Array
class OptionWarp(PyTreeNode):
"""Derived fields from Option."""
broadphase: int
broadphase_filter: int
ccd_iterations: int
@@ -159,17 +159,13 @@ class OptionWarp(PyTreeNode):
graph_conditional: bool
graph_mode: GraphMode
impratio_invsqrt: jax.Array
ls_parallel: bool
ls_parallel_min_step: float
run_collision_detection: bool
sdf_initpoints: int
sdf_iterations: int
sleep_tolerance: jax.Array
class ModelWarp(PyTreeNode):
"""Derived fields from Model."""
D_colind: np.ndarray
D_diag: np.ndarray
D_rowadr: np.ndarray
@@ -254,6 +250,7 @@ class ModelWarp(PyTreeNode):
is_sparse: bool
jnt_limited_ball_adr: np.ndarray
jnt_limited_slide_hinge_adr: np.ndarray
jtcj_max_pairs: int
light_bodyid: np.ndarray
light_targetbodyid: np.ndarray
mapD2M: np.ndarray
@@ -351,10 +348,8 @@ class ModelWarp(PyTreeNode):
wrap_site_adr: np.ndarray
wrap_site_pair_adr: np.ndarray
class DataWarp(PyTreeNode):
"""Derived fields from Data."""
M: jax.Array
actuator_moment: jax.Array
actuator_velocity: jax.Array
@@ -473,8 +468,6 @@ class DataWarp(PyTreeNode):
wrap_obj: jax.Array
wrap_xpos: jax.Array
shape = property(lambda self: self.cacc.shape)
DATA_NON_VMAP = {
'contact__dim',
'contact__dist',
@@ -503,7 +496,6 @@ DATA_NON_VMAP = {
'nworld',
}
def _to_elt(cont, _, d, axis):
return DataWarp(**{
f.name: (
@@ -749,10 +741,14 @@ _NDIM = {
'block_dim__render': 0,
'block_dim__segmented_sort': 0,
'block_dim__solve_LD_sparse_fused': 0,
'block_dim__solve_beta_accumulate': 0,
'block_dim__solve_init_search_cg': 0,
'block_dim__solve_search_update_cg': 0,
'block_dim__update_gradient_JTDAJ_dense': 0,
'block_dim__update_gradient_JTDAJ_sparse': 0,
'block_dim__update_gradient_cholesky': 0,
'block_dim__update_gradient_cholesky_blocked': 0,
'block_dim__update_gradient_grad': 0,
'body_branch_start': 1,
'body_branches': 1,
'body_conaffinity': 1,
@@ -914,6 +910,7 @@ _NDIM = {
'jnt_stiffness': 2,
'jnt_stiffnesspoly': 3,
'jnt_type': 1,
'jtcj_max_pairs': 0,
'light_active': 2,
'light_ambient': 3,
'light_attenuation': 3,
@@ -1043,8 +1040,6 @@ _NDIM = {
'opt__integrator': 0,
'opt__iterations': 0,
'opt__ls_iterations': 0,
'opt__ls_parallel': 0,
'opt__ls_parallel_min_step': 0,
'opt__ls_tolerance': 1,
'opt__magnetic': 2,
'opt__run_collision_detection': 0,
@@ -1170,8 +1165,6 @@ _NDIM = {
'integrator': 0,
'iterations': 0,
'ls_iterations': 0,
'ls_parallel': 0,
'ls_parallel_min_step': 0,
'ls_tolerance': 1,
'magnetic': 2,
'run_collision_detection': 0,
@@ -1407,10 +1400,14 @@ _BATCH_DIM = {
'block_dim__render': False,
'block_dim__segmented_sort': False,
'block_dim__solve_LD_sparse_fused': False,
'block_dim__solve_beta_accumulate': False,
'block_dim__solve_init_search_cg': False,
'block_dim__solve_search_update_cg': False,
'block_dim__update_gradient_JTDAJ_dense': False,
'block_dim__update_gradient_JTDAJ_sparse': False,
'block_dim__update_gradient_cholesky': False,
'block_dim__update_gradient_cholesky_blocked': False,
'block_dim__update_gradient_grad': False,
'body_branch_start': False,
'body_branches': False,
'body_conaffinity': False,
@@ -1572,6 +1569,7 @@ _BATCH_DIM = {
'jnt_stiffness': True,
'jnt_stiffnesspoly': True,
'jnt_type': False,
'jtcj_max_pairs': False,
'light_active': True,
'light_ambient': True,
'light_attenuation': True,
@@ -1701,8 +1699,6 @@ _BATCH_DIM = {
'opt__integrator': False,
'opt__iterations': False,
'opt__ls_iterations': False,
'opt__ls_parallel': False,
'opt__ls_parallel_min_step': False,
'opt__ls_tolerance': True,
'opt__magnetic': True,
'opt__run_collision_detection': False,
@@ -1828,8 +1824,6 @@ _BATCH_DIM = {
'integrator': False,
'iterations': False,
'ls_iterations': False,
'ls_parallel': False,
'ls_parallel_min_step': False,
'ls_tolerance': True,
'magnetic': True,
'run_collision_detection': False,