# Copyright 2026 DeepMind Technologies Limited # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. # You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # ============================================================================== """DO NOT EDIT. This file is auto-generated.""" import dataclasses import functools import jax import warp as wp from mujoco.mjx._src import types 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 _m = mjwarp.Model( **{f.name: None for f in dataclasses.fields(mjwarp.Model) if f.init} ) _d = mjwarp.Data( **{f.name: None for f in dataclasses.fields(mjwarp.Data) if f.init} ) _o = mjwarp.Option( **{f.name: None for f in dataclasses.fields(mjwarp.Option) if f.init} ) _s = mjwarp.Statistic( **{f.name: None for f in dataclasses.fields(mjwarp.Statistic) if f.init} ) _c = mjwarp.Contact( **{f.name: None for f in dataclasses.fields(mjwarp.Contact) if f.init} ) _e = mjwarp.Constraint( **{f.name: None for f in dataclasses.fields(mjwarp.Constraint) if f.init} ) _cb = mjwp_types.Callback( **{f.name: None for f in dataclasses.fields(mjwp_types.Callback) if f.init} ) @ffi.format_args_for_warp def _forward_shim( # Model nworld: int, M_elemid: wp.array2d[int], M_fullm_i: wp.array[int], M_fullm_j: wp.array[int], M_mulm_col: wp.array[int], M_mulm_madr: wp.array[int], M_mulm_rowadr: wp.array[int], M_rowadr: wp.array[int], M_rownnz: wp.array[int], M_tiles: tuple[mjwp_types.TileSet, ...], actuator_acc0: wp.array2d[float], actuator_actadr: wp.array[int], actuator_actearly: wp.array[bool], actuator_actlimited: wp.array[bool], actuator_actnum: wp.array[int], actuator_actrange: wp.array2d[wp.vec2], actuator_biasprm: wp.array2d[mjwp_types.vec10f], actuator_biastype: wp.array[int], actuator_cranklength: wp.array2d[float], actuator_ctrllimited: wp.array[bool], actuator_ctrlrange: wp.array2d[wp.vec2], actuator_delay: wp.array[float], actuator_dynprm: wp.array2d[mjwp_types.vec10f], actuator_dyntype: wp.array[int], actuator_forcelimited: wp.array[bool], actuator_forcerange: wp.array2d[wp.vec2], actuator_gainprm: wp.array2d[mjwp_types.vec10f], actuator_gaintype: wp.array[int], actuator_gear: wp.array2d[wp.spatial_vector], actuator_history: wp.array[wp.vec2i], actuator_historyadr: wp.array[int], actuator_lengthrange: wp.array2d[wp.vec2], actuator_trnid: wp.array[wp.vec2i], actuator_trntype: wp.array[int], actuator_trntype_body_adr: wp.array[int], block_dim: mjwp_types.BlockDim, body_branch_start: wp.array[int], body_branches: wp.array[int], body_dofadr: wp.array[int], body_dofnum: wp.array[int], body_fluid_ellipsoid: wp.array[bool], body_geomadr: wp.array[int], body_geomnum: wp.array[int], body_gravcomp: wp.array2d[float], body_inertia: wp.array2d[wp.vec3], body_invweight0: wp.array2d[wp.vec2], body_ipos: wp.array2d[wp.vec3], body_iquat: wp.array2d[wp.quat], body_isdofancestor: wp.array2d[int], body_jntadr: wp.array[int], body_jntnum: wp.array[int], body_mass: wp.array2d[float], body_mocapid: wp.array[int], body_parentid: wp.array[int], body_pos: wp.array2d[wp.vec3], body_quat: wp.array2d[wp.quat], body_rootid: wp.array[int], body_subtreemass: wp.array2d[float], body_tree: tuple[wp.array[int], ...], body_treeid: wp.array[int], body_weldid: wp.array[int], cam_bodyid: wp.array[int], cam_fovy: wp.array2d[float], cam_intrinsic: wp.array2d[wp.vec4], cam_mat0: wp.array2d[wp.mat33], cam_mode: wp.array[int], cam_pos: wp.array2d[wp.vec3], cam_pos0: wp.array2d[wp.vec3], cam_poscom0: wp.array2d[wp.vec3], cam_quat: wp.array2d[wp.quat], cam_resolution: wp.array[wp.vec2i], cam_sensorsize: wp.array[wp.vec2], cam_targetbodyid: wp.array[int], dof_armature: wp.array2d[float], dof_bodyid: wp.array[int], dof_damping: wp.array2d[float], dof_dampingpoly: wp.array2d[wp.vec2], dof_frictionloss: wp.array2d[float], dof_invweight0: wp.array2d[float], dof_jntid: wp.array[int], dof_length: wp.array[float], dof_parentid: wp.array[int], dof_solimp: wp.array2d[mjwp_types.vec5], dof_solref: wp.array2d[wp.vec2], dof_treeid: wp.array[int], dof_tri_col: wp.array[int], dof_tri_row: wp.array[int], eq_connect_adr: wp.array[int], eq_data: wp.array2d[mjwp_types.vec11], eq_flex_adr: wp.array[int], eq_jnt_adr: wp.array[int], eq_obj1id: wp.array[int], eq_obj2id: wp.array[int], eq_objtype: wp.array[int], eq_solimp: wp.array2d[mjwp_types.vec5], eq_solref: wp.array2d[wp.vec2], eq_ten_adr: wp.array[int], eq_type: wp.array[int], eq_wld_adr: wp.array[int], flex_bending: wp.array[float], flex_bendingadr: wp.array[int], flex_centered: wp.array[bool], flex_conaffinity: wp.array[int], flex_condim: wp.array[int], flex_contype: wp.array[int], flex_damping: wp.array[float], flex_dim: wp.array[int], flex_edge: wp.array[wp.vec2i], flex_edgeadr: wp.array[int], flex_edgeflap: wp.array[wp.vec2i], flex_edgenum: wp.array[int], flex_elem: wp.array[int], flex_elemadr: wp.array[int], flex_elemdataadr: wp.array[int], flex_elemedge: wp.array[int], flex_elemedgeadr: wp.array[int], flex_elemnum: wp.array[int], flex_friction: wp.array[wp.vec3], flex_gap: wp.array[float], flex_margin: wp.array[float], flex_priority: wp.array[int], flex_radius: wp.array[float], flex_shell: wp.array[int], flex_shelldataadr: wp.array[int], flex_shellnum: wp.array[int], flex_solimp: wp.array[mjwp_types.vec5], flex_solmix: wp.array[float], flex_solref: wp.array[wp.vec2], flex_stiffness: wp.array[float], flex_stiffnessadr: wp.array[int], flex_vert: wp.array[wp.vec3], flex_vertadr: wp.array[int], flex_vertbodyid: wp.array[int], flex_vertflexid: wp.array[int], flex_vertnum: wp.array[int], flexedge_J_colind: wp.array[int], flexedge_J_rowadr: wp.array[int], flexedge_J_rownnz: wp.array[int], flexedge_invweight0: wp.array[float], flexedge_length0: wp.array[float], geom_aabb: wp.array3d[wp.vec3], geom_bodyid: wp.array[int], geom_conaffinity: wp.array[int], geom_condim: wp.array[int], geom_contype: wp.array[int], geom_dataid: wp.array2d[int], geom_fluid: wp.array2d[float], geom_friction: wp.array2d[wp.vec3], geom_gap: wp.array2d[float], geom_group: wp.array[int], geom_margin: wp.array2d[float], geom_matid: wp.array2d[int], geom_pair_type_count: tuple[int, ...], geom_plugin_index: wp.array[int], geom_pos: wp.array2d[wp.vec3], geom_priority: wp.array[int], geom_quat: wp.array2d[wp.quat], geom_rbound: wp.array2d[float], geom_rgba: wp.array2d[wp.vec4], geom_size: wp.array2d[wp.vec3], geom_solimp: wp.array2d[mjwp_types.vec5], geom_solmix: wp.array2d[float], geom_solref: wp.array2d[wp.vec2], geom_type: wp.array[int], has_fluid: bool, has_sdf_geom: bool, hfield_adr: wp.array[int], hfield_data: wp.array[float], hfield_ncol: wp.array[int], hfield_nrow: wp.array[int], hfield_size: wp.array[wp.vec4], is_sparse: bool, jnt_actfrclimited: wp.array[bool], jnt_actfrcrange: wp.array2d[wp.vec2], jnt_actgravcomp: wp.array[int], jnt_axis: wp.array2d[wp.vec3], jnt_bodyid: wp.array[int], jnt_dofadr: wp.array[int], jnt_limited_ball_adr: wp.array[int], jnt_limited_slide_hinge_adr: wp.array[int], jnt_margin: wp.array2d[float], jnt_pos: wp.array2d[wp.vec3], jnt_qposadr: wp.array[int], jnt_range: wp.array2d[wp.vec2], jnt_solimp: wp.array2d[mjwp_types.vec5], jnt_solref: wp.array2d[wp.vec2], jnt_stiffness: wp.array2d[float], jnt_stiffnesspoly: wp.array2d[wp.vec2], jnt_type: wp.array[int], light_bodyid: wp.array[int], light_dir: wp.array2d[wp.vec3], light_dir0: wp.array2d[wp.vec3], light_mode: wp.array[int], light_pos: wp.array2d[wp.vec3], light_pos0: wp.array2d[wp.vec3], light_poscom0: wp.array2d[wp.vec3], light_targetbodyid: wp.array[int], mat_rgba: wp.array2d[wp.vec4], max_ten_J_rownnz: int, mesh_face: wp.array[wp.vec3i], mesh_faceadr: wp.array[int], mesh_graph: wp.array[int], mesh_graphadr: wp.array[int], mesh_normal: wp.array[wp.vec3], mesh_normaladr: wp.array[int], mesh_normalnum: wp.array[int], mesh_octadr: wp.array[int], mesh_polyadr: wp.array[int], mesh_polymap: wp.array[int], mesh_polymapadr: wp.array[int], mesh_polymapnum: wp.array[int], mesh_polynormal: wp.array[wp.vec3], mesh_polynum: wp.array[int], mesh_polyvert: wp.array[int], mesh_polyvertadr: wp.array[int], mesh_polyvertnum: wp.array[int], mesh_quat: wp.array[wp.quat], mesh_vert: wp.array[wp.vec3], mesh_vertadr: wp.array[int], mesh_vertnum: wp.array[int], nJten: int, na: int, nacttrnbody: int, nbody: int, nbranch: int, ncam: int, neq: int, nflex: int, nflexedge: int, nflexelem: int, nflexshelldata: int, nflexvert: int, ngeom: int, ngravcomp: int, nhistory: int, njnt: int, nlight: int, nmaxcondim: int, nmaxmeshdeg: int, nmaxpolygon: int, nmaxpyramid: int, nmeshface: int, nrangefinder: int, nsensorcollision: int, nsensorcontact: int, nsensortaxel: int, nsite: int, ntendon: int, ntree: int, nu: int, nv: int, nv_pad: int, nwrap: int, nxn_geom_pair_filtered: wp.array[wp.vec2i], nxn_pairid: wp.array[wp.vec2i], nxn_pairid_filtered: wp.array[wp.vec2i], oct_aabb: wp.array2d[wp.vec3], oct_child: wp.array[mjwp_types.vec8i], oct_coeff: wp.array[mjwp_types.vec8], pair_dim: wp.array[int], pair_friction: wp.array2d[mjwp_types.vec5], pair_gap: wp.array2d[float], pair_margin: wp.array2d[float], pair_solimp: wp.array2d[mjwp_types.vec5], pair_solref: wp.array2d[wp.vec2], pair_solreffriction: wp.array2d[wp.vec2], plugin: wp.array[int], plugin_attr: wp.array[mjwp_types.vec_pluginattr], qLD_all_updates: wp.array[wp.vec3i], qLD_level_offsets: wp.array[int], qLD_updates: tuple[wp.array[wp.vec3i], ...], qpos0: wp.array2d[float], qpos_spring: wp.array2d[float], rangefinder_sensor_adr: wp.array[int], sensor_acc_adr: wp.array[int], sensor_adr: wp.array[int], sensor_adr_to_contact_adr: wp.array[int], sensor_contact_adr: wp.array[int], sensor_cutoff: wp.array[float], sensor_datatype: wp.array[int], sensor_delay: wp.array[float], sensor_dim: wp.array[int], sensor_e_kinetic: bool, sensor_e_potential: bool, sensor_history: wp.array[wp.vec2i], sensor_historyadr: wp.array[int], sensor_interval: wp.array[wp.vec2], sensor_intprm: wp.array2d[int], sensor_limitfrc_adr: wp.array[int], sensor_limitpos_adr: wp.array[int], sensor_limitvel_adr: wp.array[int], sensor_objid: wp.array[int], sensor_objtype: wp.array[int], sensor_pos_adr: wp.array[int], sensor_rangefinder_adr: wp.array[int], sensor_rangefinder_bodyid: wp.array[int], sensor_refid: wp.array[int], sensor_reftype: wp.array[int], sensor_rne_postconstraint: bool, sensor_subtree_vel: bool, sensor_tendonactfrc_adr: wp.array[int], sensor_touch_adr: wp.array[int], sensor_type: wp.array[int], sensor_vel_adr: wp.array[int], site_bodyid: wp.array[int], site_pos: wp.array2d[wp.vec3], site_quat: wp.array2d[wp.quat], site_size: wp.array[wp.vec3], site_type: wp.array[int], taxel_sensorid: wp.array[int], taxel_vertadr: wp.array[int], ten_J_colind: wp.array[int], ten_J_rowadr: wp.array[int], ten_J_rownnz: wp.array[int], tendon_actfrclimited: wp.array[bool], tendon_actfrcrange: wp.array2d[wp.vec2], tendon_adr: wp.array[int], tendon_armature: wp.array2d[float], tendon_damping: wp.array2d[float], tendon_dampingpoly: wp.array2d[wp.vec2], tendon_frictionloss: wp.array2d[float], tendon_geom_adr: wp.array[int], tendon_invweight0: wp.array2d[float], tendon_jnt_adr: wp.array[int], tendon_length0: wp.array2d[float], tendon_lengthspring: wp.array2d[wp.vec2], tendon_limited: wp.array[int], tendon_limited_adr: wp.array[int], tendon_margin: wp.array2d[float], tendon_num: wp.array[int], tendon_range: wp.array2d[wp.vec2], tendon_site_pair_adr: wp.array[int], tendon_solimp_fri: wp.array2d[mjwp_types.vec5], tendon_solimp_lim: wp.array2d[mjwp_types.vec5], tendon_solref_fri: wp.array2d[wp.vec2], tendon_solref_lim: wp.array2d[wp.vec2], tendon_stiffness: wp.array2d[float], tendon_stiffnesspoly: wp.array2d[wp.vec2], tree_dofadr: wp.array[int], tree_dofnum: wp.array[int], tree_sleep_policy: wp.array[int], wrap_geom_adr: wp.array[int], wrap_jnt_adr: wp.array[int], wrap_objid: wp.array[int], wrap_prm: wp.array[float], wrap_pulley_scale: wp.array[float], wrap_site_pair_adr: wp.array[int], wrap_type: wp.array[int], opt__broadphase: int, opt__broadphase_filter: int, opt__ccd_iterations: int, opt__ccd_tolerance: wp.array[float], opt__cone: int, opt__contact_sensor_maxmatch: int, opt__density: wp.array[float], opt__disableflags: int, opt__enableflags: int, opt__graph_conditional: bool, opt__gravity: wp.array[wp.vec3], opt__impratio_invsqrt: wp.array[float], 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, opt__sdf_initpoints: int, opt__sdf_iterations: int, opt__solver: int, opt__timestep: wp.array[float], opt__tolerance: wp.array[float], opt__viscosity: wp.array[float], opt__wind: wp.array[wp.vec3], stat__meaninertia: wp.array[float], # Data naccdmax: int, naconmax: int, njmax: int, njmax_nnz: int, M: wp.array3d[float], act: wp.array2d[float], act_dot: wp.array2d[float], actuator_force: wp.array2d[float], actuator_length: wp.array2d[float], actuator_moment: wp.array2d[float], actuator_velocity: wp.array2d[float], body_awake: wp.array2d[int], body_awake_ind: wp.array2d[int], cacc: wp.array2d[wp.spatial_vector], cam_xmat: wp.array2d[wp.mat33], cam_xpos: wp.array2d[wp.vec3], cdof: wp.array2d[wp.spatial_vector], cdof_dot: wp.array2d[wp.spatial_vector], cfrc_ext: wp.array2d[wp.spatial_vector], cfrc_int: wp.array2d[wp.spatial_vector], cinert: wp.array2d[mjwp_types.vec10], crb: wp.array2d[mjwp_types.vec10], ctrl: wp.array2d[float], cvel: wp.array2d[wp.spatial_vector], dof_awake_ind: wp.array2d[int], dof_island: wp.array2d[int], dof_islandid: wp.array2d[int], efc_islandid: wp.array2d[int], energy: wp.array[wp.vec2], eq_active: wp.array2d[bool], flexedge_J: wp.array2d[float], flexedge_length: wp.array2d[float], flexedge_velocity: wp.array2d[float], flexvert_xpos: wp.array2d[wp.vec3], geom_xmat: wp.array2d[wp.mat33], geom_xpos: wp.array2d[wp.vec3], history: wp.array2d[float], iqacc: wp.array2d[float], iqacc_smooth: wp.array2d[float], iqfrc_constraint: wp.array2d[float], iqfrc_smooth: wp.array2d[float], island_dofadr: wp.array2d[int], island_efcadr: wp.array2d[int], island_idofadr: wp.array2d[int], island_ne: wp.array2d[int], island_nefc: wp.array2d[int], island_nf: wp.array2d[int], island_nv: wp.array2d[int], light_xdir: wp.array2d[wp.vec3], light_xpos: wp.array2d[wp.vec3], map_dof2idof: wp.array2d[int], map_efc2iefc: wp.array2d[int], map_idof2dof: wp.array2d[int], map_iefc2efc: wp.array2d[int], mocap_pos: wp.array2d[wp.vec3], mocap_quat: wp.array2d[wp.quat], moment_colind: wp.array2d[int], moment_rowadr: wp.array2d[int], moment_rownnz: wp.array2d[int], nacon: wp.array[int], nbody_awake: wp.array[int], ncollision: wp.array[int], ne: wp.array[int], nefc: wp.array[int], nf: wp.array[int], nidof: wp.array[int], nisland: wp.array[int], nl: wp.array[int], ntree_awake: wp.array[int], nv_awake: wp.array[int], qLD: wp.array3d[float], qLDiagInv: wp.array2d[float], qacc: wp.array2d[float], qacc_smooth: wp.array2d[float], qacc_warmstart: wp.array2d[float], qfrc_actuator: wp.array2d[float], qfrc_applied: wp.array2d[float], qfrc_bias: wp.array2d[float], qfrc_constraint: wp.array2d[float], qfrc_damper: wp.array2d[float], qfrc_fluid: wp.array2d[float], qfrc_gravcomp: wp.array2d[float], qfrc_passive: wp.array2d[float], qfrc_smooth: wp.array2d[float], qfrc_spring: wp.array2d[float], qpos: wp.array2d[float], qvel: wp.array2d[float], sensordata: wp.array2d[float], site_xmat: wp.array2d[wp.mat33], site_xpos: wp.array2d[wp.vec3], solver_niter: wp.array[int], subtree_angmom: wp.array2d[wp.vec3], subtree_com: wp.array2d[wp.vec3], subtree_linvel: wp.array2d[wp.vec3], ten_J: wp.array2d[float], ten_length: wp.array2d[float], ten_velocity: wp.array2d[float], ten_wrapadr: wp.array2d[int], ten_wrapnum: wp.array2d[int], time: wp.array[float], tree_asleep: wp.array2d[int], tree_awake: wp.array2d[int], tree_island: wp.array2d[int], wrap_obj: wp.array2d[wp.vec2i], wrap_xpos: wp.array2d[wp.spatial_vector], xanchor: wp.array2d[wp.vec3], xaxis: wp.array2d[wp.vec3], xfrc_applied: wp.array2d[wp.spatial_vector], ximat: wp.array2d[wp.mat33], xipos: wp.array2d[wp.vec3], xmat: wp.array2d[wp.mat33], xpos: wp.array2d[wp.vec3], xquat: wp.array2d[wp.quat], contact__dim: wp.array[int], contact__dist: wp.array[float], contact__efc_address: wp.array2d[int], contact__flex: wp.array[wp.vec2i], contact__frame: wp.array[wp.mat33], contact__friction: wp.array[mjwp_types.vec5], contact__geom: wp.array[wp.vec2i], contact__geomcollisionid: wp.array[int], contact__includemargin: wp.array[float], contact__pos: wp.array[wp.vec3], contact__solimp: wp.array[mjwp_types.vec5], contact__solref: wp.array[wp.vec2], contact__solreffriction: wp.array[wp.vec2], contact__type: wp.array[int], contact__vert: wp.array[wp.vec2i], contact__worldid: wp.array[int], efc__D: wp.array2d[float], efc__J: wp.array3d[float], efc__J_colind: wp.array3d[int], efc__J_rowadr: wp.array2d[int], efc__J_rownnz: wp.array2d[int], efc__Jqvel: wp.array2d[float], efc__Ma: wp.array2d[float], efc__aref: wp.array2d[float], efc__force: wp.array2d[float], efc__frictionloss: wp.array2d[float], efc__iD: wp.array2d[float], efc__iJ: wp.array3d[float], efc__iJ_colind: wp.array3d[int], efc__iJ_rowadr: wp.array2d[int], efc__iJ_rownnz: wp.array2d[int], efc__iaref: wp.array2d[float], efc__id: wp.array2d[int], efc__iforce: wp.array2d[float], efc__ifrictionloss: wp.array2d[float], efc__iid: wp.array2d[int], efc__island: wp.array2d[int], efc__istate: wp.array2d[int], efc__itype: wp.array2d[int], efc__margin: wp.array2d[float], efc__pos: wp.array2d[float], efc__state: wp.array2d[int], efc__type: wp.array2d[int], efc__vel: wp.array2d[float], ): _m.stat = _s _m.opt = _o _m.callback = _cb _d.efc = _e _d.contact = _c _m.M_elemid = M_elemid _m.M_fullm_i = M_fullm_i _m.M_fullm_j = M_fullm_j _m.M_mulm_col = M_mulm_col _m.M_mulm_madr = M_mulm_madr _m.M_mulm_rowadr = M_mulm_rowadr _m.M_rowadr = M_rowadr _m.M_rownnz = M_rownnz _m.M_tiles = M_tiles _m.actuator_acc0 = actuator_acc0 _m.actuator_actadr = actuator_actadr _m.actuator_actearly = actuator_actearly _m.actuator_actlimited = actuator_actlimited _m.actuator_actnum = actuator_actnum _m.actuator_actrange = actuator_actrange _m.actuator_biasprm = actuator_biasprm _m.actuator_biastype = actuator_biastype _m.actuator_cranklength = actuator_cranklength _m.actuator_ctrllimited = actuator_ctrllimited _m.actuator_ctrlrange = actuator_ctrlrange _m.actuator_delay = actuator_delay _m.actuator_dynprm = actuator_dynprm _m.actuator_dyntype = actuator_dyntype _m.actuator_forcelimited = actuator_forcelimited _m.actuator_forcerange = actuator_forcerange _m.actuator_gainprm = actuator_gainprm _m.actuator_gaintype = actuator_gaintype _m.actuator_gear = actuator_gear _m.actuator_history = actuator_history _m.actuator_historyadr = actuator_historyadr _m.actuator_lengthrange = actuator_lengthrange _m.actuator_trnid = actuator_trnid _m.actuator_trntype = actuator_trntype _m.actuator_trntype_body_adr = actuator_trntype_body_adr _m.block_dim = block_dim _m.body_branch_start = body_branch_start _m.body_branches = body_branches _m.body_dofadr = body_dofadr _m.body_dofnum = body_dofnum _m.body_fluid_ellipsoid = body_fluid_ellipsoid _m.body_geomadr = body_geomadr _m.body_geomnum = body_geomnum _m.body_gravcomp = body_gravcomp _m.body_inertia = body_inertia _m.body_invweight0 = body_invweight0 _m.body_ipos = body_ipos _m.body_iquat = body_iquat _m.body_isdofancestor = body_isdofancestor _m.body_jntadr = body_jntadr _m.body_jntnum = body_jntnum _m.body_mass = body_mass _m.body_mocapid = body_mocapid _m.body_parentid = body_parentid _m.body_pos = body_pos _m.body_quat = body_quat _m.body_rootid = body_rootid _m.body_subtreemass = body_subtreemass _m.body_tree = body_tree _m.body_treeid = body_treeid _m.body_weldid = body_weldid _m.cam_bodyid = cam_bodyid _m.cam_fovy = cam_fovy _m.cam_intrinsic = cam_intrinsic _m.cam_mat0 = cam_mat0 _m.cam_mode = cam_mode _m.cam_pos = cam_pos _m.cam_pos0 = cam_pos0 _m.cam_poscom0 = cam_poscom0 _m.cam_quat = cam_quat _m.cam_resolution = cam_resolution _m.cam_sensorsize = cam_sensorsize _m.cam_targetbodyid = cam_targetbodyid _m.dof_armature = dof_armature _m.dof_bodyid = dof_bodyid _m.dof_damping = dof_damping _m.dof_dampingpoly = dof_dampingpoly _m.dof_frictionloss = dof_frictionloss _m.dof_invweight0 = dof_invweight0 _m.dof_jntid = dof_jntid _m.dof_length = dof_length _m.dof_parentid = dof_parentid _m.dof_solimp = dof_solimp _m.dof_solref = dof_solref _m.dof_treeid = dof_treeid _m.dof_tri_col = dof_tri_col _m.dof_tri_row = dof_tri_row _m.eq_connect_adr = eq_connect_adr _m.eq_data = eq_data _m.eq_flex_adr = eq_flex_adr _m.eq_jnt_adr = eq_jnt_adr _m.eq_obj1id = eq_obj1id _m.eq_obj2id = eq_obj2id _m.eq_objtype = eq_objtype _m.eq_solimp = eq_solimp _m.eq_solref = eq_solref _m.eq_ten_adr = eq_ten_adr _m.eq_type = eq_type _m.eq_wld_adr = eq_wld_adr _m.flex_bending = flex_bending _m.flex_bendingadr = flex_bendingadr _m.flex_centered = flex_centered _m.flex_conaffinity = flex_conaffinity _m.flex_condim = flex_condim _m.flex_contype = flex_contype _m.flex_damping = flex_damping _m.flex_dim = flex_dim _m.flex_edge = flex_edge _m.flex_edgeadr = flex_edgeadr _m.flex_edgeflap = flex_edgeflap _m.flex_edgenum = flex_edgenum _m.flex_elem = flex_elem _m.flex_elemadr = flex_elemadr _m.flex_elemdataadr = flex_elemdataadr _m.flex_elemedge = flex_elemedge _m.flex_elemedgeadr = flex_elemedgeadr _m.flex_elemnum = flex_elemnum _m.flex_friction = flex_friction _m.flex_gap = flex_gap _m.flex_margin = flex_margin _m.flex_priority = flex_priority _m.flex_radius = flex_radius _m.flex_shell = flex_shell _m.flex_shelldataadr = flex_shelldataadr _m.flex_shellnum = flex_shellnum _m.flex_solimp = flex_solimp _m.flex_solmix = flex_solmix _m.flex_solref = flex_solref _m.flex_stiffness = flex_stiffness _m.flex_stiffnessadr = flex_stiffnessadr _m.flex_vert = flex_vert _m.flex_vertadr = flex_vertadr _m.flex_vertbodyid = flex_vertbodyid _m.flex_vertflexid = flex_vertflexid _m.flex_vertnum = flex_vertnum _m.flexedge_J_colind = flexedge_J_colind _m.flexedge_J_rowadr = flexedge_J_rowadr _m.flexedge_J_rownnz = flexedge_J_rownnz _m.flexedge_invweight0 = flexedge_invweight0 _m.flexedge_length0 = flexedge_length0 _m.geom_aabb = geom_aabb _m.geom_bodyid = geom_bodyid _m.geom_conaffinity = geom_conaffinity _m.geom_condim = geom_condim _m.geom_contype = geom_contype _m.geom_dataid = geom_dataid _m.geom_fluid = geom_fluid _m.geom_friction = geom_friction _m.geom_gap = geom_gap _m.geom_group = geom_group _m.geom_margin = geom_margin _m.geom_matid = geom_matid _m.geom_pair_type_count = geom_pair_type_count _m.geom_plugin_index = geom_plugin_index _m.geom_pos = geom_pos _m.geom_priority = geom_priority _m.geom_quat = geom_quat _m.geom_rbound = geom_rbound _m.geom_rgba = geom_rgba _m.geom_size = geom_size _m.geom_solimp = geom_solimp _m.geom_solmix = geom_solmix _m.geom_solref = geom_solref _m.geom_type = geom_type _m.has_fluid = has_fluid _m.has_sdf_geom = has_sdf_geom _m.hfield_adr = hfield_adr _m.hfield_data = hfield_data _m.hfield_ncol = hfield_ncol _m.hfield_nrow = hfield_nrow _m.hfield_size = hfield_size _m.is_sparse = is_sparse _m.jnt_actfrclimited = jnt_actfrclimited _m.jnt_actfrcrange = jnt_actfrcrange _m.jnt_actgravcomp = jnt_actgravcomp _m.jnt_axis = jnt_axis _m.jnt_bodyid = jnt_bodyid _m.jnt_dofadr = jnt_dofadr _m.jnt_limited_ball_adr = jnt_limited_ball_adr _m.jnt_limited_slide_hinge_adr = jnt_limited_slide_hinge_adr _m.jnt_margin = jnt_margin _m.jnt_pos = jnt_pos _m.jnt_qposadr = jnt_qposadr _m.jnt_range = jnt_range _m.jnt_solimp = jnt_solimp _m.jnt_solref = jnt_solref _m.jnt_stiffness = jnt_stiffness _m.jnt_stiffnesspoly = jnt_stiffnesspoly _m.jnt_type = jnt_type _m.light_bodyid = light_bodyid _m.light_dir = light_dir _m.light_dir0 = light_dir0 _m.light_mode = light_mode _m.light_pos = light_pos _m.light_pos0 = light_pos0 _m.light_poscom0 = light_poscom0 _m.light_targetbodyid = light_targetbodyid _m.mat_rgba = mat_rgba _m.max_ten_J_rownnz = max_ten_J_rownnz _m.mesh_face = mesh_face _m.mesh_faceadr = mesh_faceadr _m.mesh_graph = mesh_graph _m.mesh_graphadr = mesh_graphadr _m.mesh_normal = mesh_normal _m.mesh_normaladr = mesh_normaladr _m.mesh_normalnum = mesh_normalnum _m.mesh_octadr = mesh_octadr _m.mesh_polyadr = mesh_polyadr _m.mesh_polymap = mesh_polymap _m.mesh_polymapadr = mesh_polymapadr _m.mesh_polymapnum = mesh_polymapnum _m.mesh_polynormal = mesh_polynormal _m.mesh_polynum = mesh_polynum _m.mesh_polyvert = mesh_polyvert _m.mesh_polyvertadr = mesh_polyvertadr _m.mesh_polyvertnum = mesh_polyvertnum _m.mesh_quat = mesh_quat _m.mesh_vert = mesh_vert _m.mesh_vertadr = mesh_vertadr _m.mesh_vertnum = mesh_vertnum _m.nJten = nJten _m.na = na _m.nacttrnbody = nacttrnbody _m.nbody = nbody _m.nbranch = nbranch _m.ncam = ncam _m.neq = neq _m.nflex = nflex _m.nflexedge = nflexedge _m.nflexelem = nflexelem _m.nflexshelldata = nflexshelldata _m.nflexvert = nflexvert _m.ngeom = ngeom _m.ngravcomp = ngravcomp _m.nhistory = nhistory _m.njnt = njnt _m.nlight = nlight _m.nmaxcondim = nmaxcondim _m.nmaxmeshdeg = nmaxmeshdeg _m.nmaxpolygon = nmaxpolygon _m.nmaxpyramid = nmaxpyramid _m.nmeshface = nmeshface _m.nrangefinder = nrangefinder _m.nsensorcollision = nsensorcollision _m.nsensorcontact = nsensorcontact _m.nsensortaxel = nsensortaxel _m.nsite = nsite _m.ntendon = ntendon _m.ntree = ntree _m.nu = nu _m.nv = nv _m.nv_pad = nv_pad _m.nwrap = nwrap _m.nxn_geom_pair_filtered = nxn_geom_pair_filtered _m.nxn_pairid = nxn_pairid _m.nxn_pairid_filtered = nxn_pairid_filtered _m.oct_aabb = oct_aabb _m.oct_child = oct_child _m.oct_coeff = oct_coeff _m.opt.broadphase = opt__broadphase _m.opt.broadphase_filter = opt__broadphase_filter _m.opt.ccd_iterations = opt__ccd_iterations _m.opt.ccd_tolerance = opt__ccd_tolerance _m.opt.cone = opt__cone _m.opt.contact_sensor_maxmatch = opt__contact_sensor_maxmatch _m.opt.density = opt__density _m.opt.disableflags = opt__disableflags _m.opt.enableflags = opt__enableflags _m.opt.graph_conditional = opt__graph_conditional _m.opt.gravity = opt__gravity _m.opt.impratio_invsqrt = opt__impratio_invsqrt _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 _m.opt.sdf_initpoints = opt__sdf_initpoints _m.opt.sdf_iterations = opt__sdf_iterations _m.opt.solver = opt__solver _m.opt.timestep = opt__timestep _m.opt.tolerance = opt__tolerance _m.opt.viscosity = opt__viscosity _m.opt.wind = opt__wind _m.pair_dim = pair_dim _m.pair_friction = pair_friction _m.pair_gap = pair_gap _m.pair_margin = pair_margin _m.pair_solimp = pair_solimp _m.pair_solref = pair_solref _m.pair_solreffriction = pair_solreffriction _m.plugin = plugin _m.plugin_attr = plugin_attr _m.qLD_all_updates = qLD_all_updates _m.qLD_level_offsets = qLD_level_offsets _m.qLD_updates = qLD_updates _m.qpos0 = qpos0 _m.qpos_spring = qpos_spring _m.rangefinder_sensor_adr = rangefinder_sensor_adr _m.sensor_acc_adr = sensor_acc_adr _m.sensor_adr = sensor_adr _m.sensor_adr_to_contact_adr = sensor_adr_to_contact_adr _m.sensor_contact_adr = sensor_contact_adr _m.sensor_cutoff = sensor_cutoff _m.sensor_datatype = sensor_datatype _m.sensor_delay = sensor_delay _m.sensor_dim = sensor_dim _m.sensor_e_kinetic = sensor_e_kinetic _m.sensor_e_potential = sensor_e_potential _m.sensor_history = sensor_history _m.sensor_historyadr = sensor_historyadr _m.sensor_interval = sensor_interval _m.sensor_intprm = sensor_intprm _m.sensor_limitfrc_adr = sensor_limitfrc_adr _m.sensor_limitpos_adr = sensor_limitpos_adr _m.sensor_limitvel_adr = sensor_limitvel_adr _m.sensor_objid = sensor_objid _m.sensor_objtype = sensor_objtype _m.sensor_pos_adr = sensor_pos_adr _m.sensor_rangefinder_adr = sensor_rangefinder_adr _m.sensor_rangefinder_bodyid = sensor_rangefinder_bodyid _m.sensor_refid = sensor_refid _m.sensor_reftype = sensor_reftype _m.sensor_rne_postconstraint = sensor_rne_postconstraint _m.sensor_subtree_vel = sensor_subtree_vel _m.sensor_tendonactfrc_adr = sensor_tendonactfrc_adr _m.sensor_touch_adr = sensor_touch_adr _m.sensor_type = sensor_type _m.sensor_vel_adr = sensor_vel_adr _m.site_bodyid = site_bodyid _m.site_pos = site_pos _m.site_quat = site_quat _m.site_size = site_size _m.site_type = site_type _m.stat.meaninertia = stat__meaninertia _m.taxel_sensorid = taxel_sensorid _m.taxel_vertadr = taxel_vertadr _m.ten_J_colind = ten_J_colind _m.ten_J_rowadr = ten_J_rowadr _m.ten_J_rownnz = ten_J_rownnz _m.tendon_actfrclimited = tendon_actfrclimited _m.tendon_actfrcrange = tendon_actfrcrange _m.tendon_adr = tendon_adr _m.tendon_armature = tendon_armature _m.tendon_damping = tendon_damping _m.tendon_dampingpoly = tendon_dampingpoly _m.tendon_frictionloss = tendon_frictionloss _m.tendon_geom_adr = tendon_geom_adr _m.tendon_invweight0 = tendon_invweight0 _m.tendon_jnt_adr = tendon_jnt_adr _m.tendon_length0 = tendon_length0 _m.tendon_lengthspring = tendon_lengthspring _m.tendon_limited = tendon_limited _m.tendon_limited_adr = tendon_limited_adr _m.tendon_margin = tendon_margin _m.tendon_num = tendon_num _m.tendon_range = tendon_range _m.tendon_site_pair_adr = tendon_site_pair_adr _m.tendon_solimp_fri = tendon_solimp_fri _m.tendon_solimp_lim = tendon_solimp_lim _m.tendon_solref_fri = tendon_solref_fri _m.tendon_solref_lim = tendon_solref_lim _m.tendon_stiffness = tendon_stiffness _m.tendon_stiffnesspoly = tendon_stiffnesspoly _m.tree_dofadr = tree_dofadr _m.tree_dofnum = tree_dofnum _m.tree_sleep_policy = tree_sleep_policy _m.wrap_geom_adr = wrap_geom_adr _m.wrap_jnt_adr = wrap_jnt_adr _m.wrap_objid = wrap_objid _m.wrap_prm = wrap_prm _m.wrap_pulley_scale = wrap_pulley_scale _m.wrap_site_pair_adr = wrap_site_pair_adr _m.wrap_type = wrap_type _d.M = M _d.act = act _d.act_dot = act_dot _d.actuator_force = actuator_force _d.actuator_length = actuator_length _d.actuator_moment = actuator_moment _d.actuator_velocity = actuator_velocity _d.body_awake = body_awake _d.body_awake_ind = body_awake_ind _d.cacc = cacc _d.cam_xmat = cam_xmat _d.cam_xpos = cam_xpos _d.cdof = cdof _d.cdof_dot = cdof_dot _d.cfrc_ext = cfrc_ext _d.cfrc_int = cfrc_int _d.cinert = cinert _d.contact.dim = contact__dim _d.contact.dist = contact__dist _d.contact.efc_address = contact__efc_address _d.contact.flex = contact__flex _d.contact.frame = contact__frame _d.contact.friction = contact__friction _d.contact.geom = contact__geom _d.contact.geomcollisionid = contact__geomcollisionid _d.contact.includemargin = contact__includemargin _d.contact.pos = contact__pos _d.contact.solimp = contact__solimp _d.contact.solref = contact__solref _d.contact.solreffriction = contact__solreffriction _d.contact.type = contact__type _d.contact.vert = contact__vert _d.contact.worldid = contact__worldid _d.crb = crb _d.ctrl = ctrl _d.cvel = cvel _d.dof_awake_ind = dof_awake_ind _d.dof_island = dof_island _d.dof_islandid = dof_islandid _d.efc.D = efc__D _d.efc.J = efc__J _d.efc.J_colind = efc__J_colind _d.efc.J_rowadr = efc__J_rowadr _d.efc.J_rownnz = efc__J_rownnz _d.efc.Jqvel = efc__Jqvel _d.efc.Ma = efc__Ma _d.efc.aref = efc__aref _d.efc.force = efc__force _d.efc.frictionloss = efc__frictionloss _d.efc.iD = efc__iD _d.efc.iJ = efc__iJ _d.efc.iJ_colind = efc__iJ_colind _d.efc.iJ_rowadr = efc__iJ_rowadr _d.efc.iJ_rownnz = efc__iJ_rownnz _d.efc.iaref = efc__iaref _d.efc.id = efc__id _d.efc.iforce = efc__iforce _d.efc.ifrictionloss = efc__ifrictionloss _d.efc.iid = efc__iid _d.efc.island = efc__island _d.efc.istate = efc__istate _d.efc.itype = efc__itype _d.efc.margin = efc__margin _d.efc.pos = efc__pos _d.efc.state = efc__state _d.efc.type = efc__type _d.efc.vel = efc__vel _d.efc_islandid = efc_islandid _d.energy = energy _d.eq_active = eq_active _d.flexedge_J = flexedge_J _d.flexedge_length = flexedge_length _d.flexedge_velocity = flexedge_velocity _d.flexvert_xpos = flexvert_xpos _d.geom_xmat = geom_xmat _d.geom_xpos = geom_xpos _d.history = history _d.iqacc = iqacc _d.iqacc_smooth = iqacc_smooth _d.iqfrc_constraint = iqfrc_constraint _d.iqfrc_smooth = iqfrc_smooth _d.island_dofadr = island_dofadr _d.island_efcadr = island_efcadr _d.island_idofadr = island_idofadr _d.island_ne = island_ne _d.island_nefc = island_nefc _d.island_nf = island_nf _d.island_nv = island_nv _d.light_xdir = light_xdir _d.light_xpos = light_xpos _d.map_dof2idof = map_dof2idof _d.map_efc2iefc = map_efc2iefc _d.map_idof2dof = map_idof2dof _d.map_iefc2efc = map_iefc2efc _d.mocap_pos = mocap_pos _d.mocap_quat = mocap_quat _d.moment_colind = moment_colind _d.moment_rowadr = moment_rowadr _d.moment_rownnz = moment_rownnz _d.naccdmax = naccdmax _d.nacon = nacon _d.naconmax = naconmax _d.nbody_awake = nbody_awake _d.ncollision = ncollision _d.ne = ne _d.nefc = nefc _d.nf = nf _d.nidof = nidof _d.nisland = nisland _d.njmax = njmax _d.njmax_nnz = njmax_nnz _d.nl = nl _d.ntree_awake = ntree_awake _d.nv_awake = nv_awake _d.qLD = qLD _d.qLDiagInv = qLDiagInv _d.qacc = qacc _d.qacc_smooth = qacc_smooth _d.qacc_warmstart = qacc_warmstart _d.qfrc_actuator = qfrc_actuator _d.qfrc_applied = qfrc_applied _d.qfrc_bias = qfrc_bias _d.qfrc_constraint = qfrc_constraint _d.qfrc_damper = qfrc_damper _d.qfrc_fluid = qfrc_fluid _d.qfrc_gravcomp = qfrc_gravcomp _d.qfrc_passive = qfrc_passive _d.qfrc_smooth = qfrc_smooth _d.qfrc_spring = qfrc_spring _d.qpos = qpos _d.qvel = qvel _d.sensordata = sensordata _d.site_xmat = site_xmat _d.site_xpos = site_xpos _d.solver_niter = solver_niter _d.subtree_angmom = subtree_angmom _d.subtree_com = subtree_com _d.subtree_linvel = subtree_linvel _d.ten_J = ten_J _d.ten_length = ten_length _d.ten_velocity = ten_velocity _d.ten_wrapadr = ten_wrapadr _d.ten_wrapnum = ten_wrapnum _d.time = time _d.tree_asleep = tree_asleep _d.tree_awake = tree_awake _d.tree_island = tree_island _d.wrap_obj = wrap_obj _d.wrap_xpos = wrap_xpos _d.xanchor = xanchor _d.xaxis = xaxis _d.xfrc_applied = xfrc_applied _d.ximat = ximat _d.xipos = xipos _d.xmat = xmat _d.xpos = xpos _d.xquat = xquat _d.nworld = nworld mjwarp.forward(_m, _d) def _forward_jax_impl(m: types.Model, d: types.Data): output_dims = { 'M': d._impl.M.shape, 'act_dot': d.act_dot.shape, 'actuator_force': d.actuator_force.shape, 'actuator_length': d.actuator_length.shape, 'actuator_moment': d._impl.actuator_moment.shape, 'actuator_velocity': d._impl.actuator_velocity.shape, 'body_awake': d._impl.body_awake.shape, 'body_awake_ind': d._impl.body_awake_ind.shape, 'cacc': d._impl.cacc.shape, 'cam_xmat': d.cam_xmat.shape, 'cam_xpos': d.cam_xpos.shape, 'cdof': d.cdof.shape, 'cdof_dot': d.cdof_dot.shape, 'cfrc_ext': d._impl.cfrc_ext.shape, 'cfrc_int': d._impl.cfrc_int.shape, 'cinert': d._impl.cinert.shape, 'crb': d._impl.crb.shape, 'cvel': d.cvel.shape, 'dof_awake_ind': d._impl.dof_awake_ind.shape, 'dof_island': d._impl.dof_island.shape, 'dof_islandid': d._impl.dof_islandid.shape, 'efc_islandid': d._impl.efc_islandid.shape, 'energy': d._impl.energy.shape, 'flexedge_J': d._impl.flexedge_J.shape, 'flexedge_length': d._impl.flexedge_length.shape, 'flexedge_velocity': d._impl.flexedge_velocity.shape, 'flexvert_xpos': d._impl.flexvert_xpos.shape, 'geom_xmat': d.geom_xmat.shape, 'geom_xpos': d.geom_xpos.shape, 'history': d.history.shape, 'iqacc': d._impl.iqacc.shape, 'iqacc_smooth': d._impl.iqacc_smooth.shape, 'iqfrc_constraint': d._impl.iqfrc_constraint.shape, 'iqfrc_smooth': d._impl.iqfrc_smooth.shape, 'island_dofadr': d._impl.island_dofadr.shape, 'island_efcadr': d._impl.island_efcadr.shape, 'island_idofadr': d._impl.island_idofadr.shape, 'island_ne': d._impl.island_ne.shape, 'island_nefc': d._impl.island_nefc.shape, 'island_nf': d._impl.island_nf.shape, 'island_nv': d._impl.island_nv.shape, 'light_xdir': d._impl.light_xdir.shape, 'light_xpos': d._impl.light_xpos.shape, 'map_dof2idof': d._impl.map_dof2idof.shape, 'map_efc2iefc': d._impl.map_efc2iefc.shape, 'map_idof2dof': d._impl.map_idof2dof.shape, 'map_iefc2efc': d._impl.map_iefc2efc.shape, 'moment_colind': d._impl.moment_colind.shape, 'moment_rowadr': d._impl.moment_rowadr.shape, 'moment_rownnz': d._impl.moment_rownnz.shape, 'nacon': d._impl.nacon.shape, 'nbody_awake': d._impl.nbody_awake.shape, 'ncollision': d._impl.ncollision.shape, 'ne': d._impl.ne.shape, 'nefc': d._impl.nefc.shape, 'nf': d._impl.nf.shape, 'nidof': d._impl.nidof.shape, 'nisland': d._impl.nisland.shape, 'nl': d._impl.nl.shape, 'ntree_awake': d._impl.ntree_awake.shape, 'nv_awake': d._impl.nv_awake.shape, 'qLD': d._impl.qLD.shape, 'qLDiagInv': d._impl.qLDiagInv.shape, 'qacc': d.qacc.shape, 'qacc_smooth': d.qacc_smooth.shape, 'qfrc_actuator': d.qfrc_actuator.shape, 'qfrc_bias': d.qfrc_bias.shape, 'qfrc_constraint': d.qfrc_constraint.shape, 'qfrc_damper': d._impl.qfrc_damper.shape, 'qfrc_fluid': d.qfrc_fluid.shape, 'qfrc_gravcomp': d.qfrc_gravcomp.shape, 'qfrc_passive': d.qfrc_passive.shape, 'qfrc_smooth': d.qfrc_smooth.shape, 'qfrc_spring': d._impl.qfrc_spring.shape, 'qvel': d.qvel.shape, 'sensordata': d.sensordata.shape, 'site_xmat': d.site_xmat.shape, 'site_xpos': d.site_xpos.shape, 'solver_niter': d._impl.solver_niter.shape, 'subtree_angmom': d._impl.subtree_angmom.shape, 'subtree_com': d.subtree_com.shape, 'subtree_linvel': d._impl.subtree_linvel.shape, 'ten_J': d._impl.ten_J.shape, 'ten_length': d.ten_length.shape, 'ten_velocity': d._impl.ten_velocity.shape, 'ten_wrapadr': d._impl.ten_wrapadr.shape, 'ten_wrapnum': d._impl.ten_wrapnum.shape, 'tree_asleep': d._impl.tree_asleep.shape, 'tree_awake': d._impl.tree_awake.shape, 'tree_island': d._impl.tree_island.shape, 'wrap_obj': d._impl.wrap_obj.shape, 'wrap_xpos': d._impl.wrap_xpos.shape, 'xanchor': d.xanchor.shape, 'xaxis': d.xaxis.shape, 'ximat': d.ximat.shape, 'xipos': d.xipos.shape, 'xmat': d.xmat.shape, 'xpos': d.xpos.shape, 'xquat': d.xquat.shape, 'contact__dim': d._impl.contact__dim.shape, 'contact__dist': d._impl.contact__dist.shape, 'contact__efc_address': d._impl.contact__efc_address.shape, 'contact__flex': d._impl.contact__flex.shape, 'contact__frame': d._impl.contact__frame.shape, 'contact__friction': d._impl.contact__friction.shape, 'contact__geom': d._impl.contact__geom.shape, 'contact__geomcollisionid': d._impl.contact__geomcollisionid.shape, 'contact__includemargin': d._impl.contact__includemargin.shape, 'contact__pos': d._impl.contact__pos.shape, 'contact__solimp': d._impl.contact__solimp.shape, 'contact__solref': d._impl.contact__solref.shape, 'contact__solreffriction': d._impl.contact__solreffriction.shape, 'contact__type': d._impl.contact__type.shape, 'contact__vert': d._impl.contact__vert.shape, 'contact__worldid': d._impl.contact__worldid.shape, 'efc__D': d._impl.efc__D.shape, 'efc__J': d._impl.efc__J.shape, 'efc__J_colind': d._impl.efc__J_colind.shape, 'efc__J_rowadr': d._impl.efc__J_rowadr.shape, 'efc__J_rownnz': d._impl.efc__J_rownnz.shape, 'efc__Jqvel': d._impl.efc__Jqvel.shape, 'efc__Ma': d._impl.efc__Ma.shape, 'efc__aref': d._impl.efc__aref.shape, 'efc__force': d._impl.efc__force.shape, 'efc__frictionloss': d._impl.efc__frictionloss.shape, 'efc__iD': d._impl.efc__iD.shape, 'efc__iJ': d._impl.efc__iJ.shape, 'efc__iJ_colind': d._impl.efc__iJ_colind.shape, 'efc__iJ_rowadr': d._impl.efc__iJ_rowadr.shape, 'efc__iJ_rownnz': d._impl.efc__iJ_rownnz.shape, 'efc__iaref': d._impl.efc__iaref.shape, 'efc__id': d._impl.efc__id.shape, 'efc__iforce': d._impl.efc__iforce.shape, 'efc__ifrictionloss': d._impl.efc__ifrictionloss.shape, 'efc__iid': d._impl.efc__iid.shape, 'efc__island': d._impl.efc__island.shape, 'efc__istate': d._impl.efc__istate.shape, 'efc__itype': d._impl.efc__itype.shape, 'efc__margin': d._impl.efc__margin.shape, 'efc__pos': d._impl.efc__pos.shape, 'efc__state': d._impl.efc__state.shape, 'efc__type': d._impl.efc__type.shape, 'efc__vel': d._impl.efc__vel.shape, } jf = ffi.jax_callable_variadic_tuple( _forward_shim, num_outputs=143, output_dims=output_dims, vmap_method=None, in_out_argnames=set([ 'M', 'act_dot', 'actuator_force', 'actuator_length', 'actuator_moment', 'actuator_velocity', 'body_awake', 'body_awake_ind', 'cacc', 'cam_xmat', 'cam_xpos', 'cdof', 'cdof_dot', 'cfrc_ext', 'cfrc_int', 'cinert', 'crb', 'cvel', 'dof_awake_ind', 'dof_island', 'dof_islandid', 'efc_islandid', 'energy', 'flexedge_J', 'flexedge_length', 'flexedge_velocity', 'flexvert_xpos', 'geom_xmat', 'geom_xpos', 'history', 'iqacc', 'iqacc_smooth', 'iqfrc_constraint', 'iqfrc_smooth', 'island_dofadr', 'island_efcadr', 'island_idofadr', 'island_ne', 'island_nefc', 'island_nf', 'island_nv', 'light_xdir', 'light_xpos', 'map_dof2idof', 'map_efc2iefc', 'map_idof2dof', 'map_iefc2efc', 'moment_colind', 'moment_rowadr', 'moment_rownnz', 'nacon', 'nbody_awake', 'ncollision', 'ne', 'nefc', 'nf', 'nidof', 'nisland', 'nl', 'ntree_awake', 'nv_awake', 'qLD', 'qLDiagInv', 'qacc', 'qacc_smooth', 'qfrc_actuator', 'qfrc_bias', 'qfrc_constraint', 'qfrc_damper', 'qfrc_fluid', 'qfrc_gravcomp', 'qfrc_passive', 'qfrc_smooth', 'qfrc_spring', 'qvel', 'sensordata', 'site_xmat', 'site_xpos', 'solver_niter', 'subtree_angmom', 'subtree_com', 'subtree_linvel', 'ten_J', 'ten_length', 'ten_velocity', 'ten_wrapadr', 'ten_wrapnum', 'tree_asleep', 'tree_awake', 'tree_island', 'wrap_obj', 'wrap_xpos', 'xanchor', 'xaxis', 'ximat', 'xipos', 'xmat', 'xpos', 'xquat', 'contact__dim', 'contact__dist', 'contact__efc_address', 'contact__flex', 'contact__frame', 'contact__friction', 'contact__geom', 'contact__geomcollisionid', 'contact__includemargin', 'contact__pos', 'contact__solimp', 'contact__solref', 'contact__solreffriction', 'contact__type', 'contact__vert', 'contact__worldid', 'efc__D', 'efc__J', 'efc__J_colind', 'efc__J_rowadr', 'efc__J_rownnz', 'efc__Jqvel', 'efc__Ma', 'efc__aref', 'efc__force', 'efc__frictionloss', 'efc__iD', 'efc__iJ', 'efc__iJ_colind', 'efc__iJ_rowadr', 'efc__iJ_rownnz', 'efc__iaref', 'efc__id', 'efc__iforce', 'efc__ifrictionloss', 'efc__iid', 'efc__island', 'efc__istate', 'efc__itype', 'efc__margin', 'efc__pos', 'efc__state', 'efc__type', 'efc__vel', ]), stage_in_argnames=set([ 'act', 'act_dot', 'actuator_acc0', 'actuator_actrange', 'actuator_biasprm', 'actuator_cranklength', 'actuator_ctrlrange', 'actuator_dynprm', 'actuator_force', 'actuator_forcerange', 'actuator_gainprm', 'actuator_gear', 'actuator_length', 'actuator_lengthrange', 'body_gravcomp', 'body_inertia', 'body_invweight0', 'body_ipos', 'body_iquat', 'body_mass', 'body_pos', 'body_quat', 'body_subtreemass', 'cam_fovy', 'cam_intrinsic', 'cam_mat0', 'cam_pos', 'cam_pos0', 'cam_poscom0', 'cam_quat', 'cam_xmat', 'cam_xpos', 'cdof', 'cdof_dot', 'ctrl', 'cvel', 'dof_armature', 'dof_damping', 'dof_dampingpoly', 'dof_frictionloss', 'dof_invweight0', 'dof_solimp', 'dof_solref', 'eq_active', 'eq_data', 'eq_solimp', 'eq_solref', 'geom_aabb', 'geom_friction', 'geom_gap', 'geom_margin', 'geom_matid', 'geom_pos', 'geom_quat', 'geom_rbound', 'geom_rgba', 'geom_size', 'geom_solimp', 'geom_solmix', 'geom_solref', 'geom_xmat', 'geom_xpos', 'hfield_data', 'history', 'jnt_actfrcrange', 'jnt_axis', 'jnt_margin', 'jnt_pos', 'jnt_range', 'jnt_solimp', 'jnt_solref', 'jnt_stiffness', 'jnt_stiffnesspoly', 'light_dir', 'light_dir0', 'light_pos', 'light_pos0', 'light_poscom0', 'mat_rgba', 'mocap_pos', 'mocap_quat', 'opt__density', 'opt__gravity', 'opt__ls_tolerance', 'opt__magnetic', 'opt__timestep', 'opt__tolerance', 'opt__viscosity', 'opt__wind', 'pair_friction', 'pair_gap', 'pair_margin', 'pair_solimp', 'pair_solref', 'pair_solreffriction', 'qacc', 'qacc_smooth', 'qacc_warmstart', 'qfrc_actuator', 'qfrc_applied', 'qfrc_bias', 'qfrc_constraint', 'qfrc_fluid', 'qfrc_gravcomp', 'qfrc_passive', 'qfrc_smooth', 'qpos', 'qpos0', 'qpos_spring', 'qvel', 'sensordata', 'site_pos', 'site_quat', 'site_xmat', 'site_xpos', 'subtree_com', 'ten_length', 'tendon_actfrcrange', 'tendon_armature', 'tendon_damping', 'tendon_dampingpoly', 'tendon_frictionloss', 'tendon_invweight0', 'tendon_length0', 'tendon_lengthspring', 'tendon_margin', 'tendon_range', 'tendon_solimp_fri', 'tendon_solimp_lim', 'tendon_solref_fri', 'tendon_solref_lim', 'tendon_stiffness', 'tendon_stiffnesspoly', 'time', 'xanchor', 'xaxis', 'xfrc_applied', 'ximat', 'xipos', 'xmat', 'xpos', 'xquat', ]), stage_out_argnames=set([ 'act_dot', 'actuator_force', 'actuator_length', 'cam_xmat', 'cam_xpos', 'cdof', 'cdof_dot', 'cvel', 'geom_xmat', 'geom_xpos', 'history', 'qacc', 'qacc_smooth', 'qfrc_actuator', 'qfrc_bias', 'qfrc_constraint', 'qfrc_fluid', 'qfrc_gravcomp', 'qfrc_passive', 'qfrc_smooth', 'qvel', 'sensordata', 'site_xmat', 'site_xpos', 'subtree_com', 'ten_length', 'xanchor', 'xaxis', 'ximat', 'xipos', 'xmat', 'xpos', 'xquat', ]), graph_mode=m.opt._impl.graph_mode, has_side_effect=False, ) out = jf( d.qpos.shape[0], m._impl.M_elemid, m._impl.M_fullm_i, m._impl.M_fullm_j, m._impl.M_mulm_col, m._impl.M_mulm_madr, m._impl.M_mulm_rowadr, m._impl.M_rowadr, m._impl.M_rownnz, m._impl.M_tiles, m.actuator_acc0, m.actuator_actadr, m.actuator_actearly, m.actuator_actlimited, m.actuator_actnum, m.actuator_actrange, m.actuator_biasprm, m.actuator_biastype, m.actuator_cranklength, m.actuator_ctrllimited, m.actuator_ctrlrange, m._impl.actuator_delay, m.actuator_dynprm, m.actuator_dyntype, m.actuator_forcelimited, m.actuator_forcerange, m.actuator_gainprm, m.actuator_gaintype, m.actuator_gear, m._impl.actuator_history, m._impl.actuator_historyadr, m.actuator_lengthrange, m.actuator_trnid, m.actuator_trntype, m._impl.actuator_trntype_body_adr, m._impl.block_dim, m._impl.body_branch_start, m._impl.body_branches, m.body_dofadr, m.body_dofnum, m._impl.body_fluid_ellipsoid, m.body_geomadr, m.body_geomnum, m.body_gravcomp, m.body_inertia, m.body_invweight0, m.body_ipos, m.body_iquat, m._impl.body_isdofancestor, m.body_jntadr, m.body_jntnum, m.body_mass, m.body_mocapid, m.body_parentid, m.body_pos, m.body_quat, m.body_rootid, m.body_subtreemass, m._impl.body_tree, m.body_treeid, m.body_weldid, m.cam_bodyid, m.cam_fovy, m.cam_intrinsic, m.cam_mat0, m.cam_mode, m.cam_pos, m.cam_pos0, m.cam_poscom0, m.cam_quat, m.cam_resolution, m.cam_sensorsize, m.cam_targetbodyid, m.dof_armature, m.dof_bodyid, m.dof_damping, m.dof_dampingpoly, m.dof_frictionloss, m.dof_invweight0, m.dof_jntid, m._impl.dof_length, m.dof_parentid, m.dof_solimp, m.dof_solref, m.dof_treeid, m._impl.dof_tri_col, m._impl.dof_tri_row, m._impl.eq_connect_adr, m.eq_data, m._impl.eq_flex_adr, m._impl.eq_jnt_adr, m.eq_obj1id, m.eq_obj2id, m.eq_objtype, m.eq_solimp, m.eq_solref, m._impl.eq_ten_adr, m.eq_type, m._impl.eq_wld_adr, m._impl.flex_bending, m._impl.flex_bendingadr, m._impl.flex_centered, m._impl.flex_conaffinity, m._impl.flex_condim, m._impl.flex_contype, m._impl.flex_damping, m._impl.flex_dim, m._impl.flex_edge, m._impl.flex_edgeadr, m._impl.flex_edgeflap, m._impl.flex_edgenum, m._impl.flex_elem, m._impl.flex_elemadr, m._impl.flex_elemdataadr, m._impl.flex_elemedge, m._impl.flex_elemedgeadr, m._impl.flex_elemnum, m._impl.flex_friction, m._impl.flex_gap, m._impl.flex_margin, m._impl.flex_priority, m._impl.flex_radius, m._impl.flex_shell, m._impl.flex_shelldataadr, m._impl.flex_shellnum, m._impl.flex_solimp, m._impl.flex_solmix, m._impl.flex_solref, m._impl.flex_stiffness, m._impl.flex_stiffnessadr, m._impl.flex_vert, m.flex_vertadr, m._impl.flex_vertbodyid, m._impl.flex_vertflexid, m.flex_vertnum, m._impl.flexedge_J_colind, m._impl.flexedge_J_rowadr, m._impl.flexedge_J_rownnz, m._impl.flexedge_invweight0, m._impl.flexedge_length0, m.geom_aabb, m.geom_bodyid, m.geom_conaffinity, m.geom_condim, m.geom_contype, jax.numpy.expand_dims(m.geom_dataid, 0), m.geom_fluid, m.geom_friction, m.geom_gap, m.geom_group, m.geom_margin, m.geom_matid, m._impl.geom_pair_type_count, m._impl.geom_plugin_index, m.geom_pos, m.geom_priority, m.geom_quat, m.geom_rbound, m.geom_rgba, m.geom_size, m.geom_solimp, m.geom_solmix, m.geom_solref, m.geom_type, m._impl.has_fluid, m._impl.has_sdf_geom, m.hfield_adr, m.hfield_data, m.hfield_ncol, m.hfield_nrow, m.hfield_size, m._impl.is_sparse, m.jnt_actfrclimited, m.jnt_actfrcrange, m.jnt_actgravcomp, m.jnt_axis, m.jnt_bodyid, m.jnt_dofadr, m._impl.jnt_limited_ball_adr, m._impl.jnt_limited_slide_hinge_adr, m.jnt_margin, m.jnt_pos, m.jnt_qposadr, m.jnt_range, m.jnt_solimp, m.jnt_solref, m.jnt_stiffness, m.jnt_stiffnesspoly, m.jnt_type, m._impl.light_bodyid, m.light_dir, m.light_dir0, m.light_mode, m.light_pos, m.light_pos0, m.light_poscom0, m._impl.light_targetbodyid, m.mat_rgba, m._impl.max_ten_J_rownnz, m.mesh_face, m.mesh_faceadr, m.mesh_graph, m.mesh_graphadr, m.mesh_normal, m.mesh_normaladr, m.mesh_normalnum, m.mesh_octadr, m._impl.mesh_polyadr, m._impl.mesh_polymap, m._impl.mesh_polymapadr, m._impl.mesh_polymapnum, m._impl.mesh_polynormal, m._impl.mesh_polynum, m._impl.mesh_polyvert, m._impl.mesh_polyvertadr, m._impl.mesh_polyvertnum, m.mesh_quat, m.mesh_vert, m.mesh_vertadr, m.mesh_vertnum, m.nJten, m.na, m._impl.nacttrnbody, m.nbody, m._impl.nbranch, m.ncam, m.neq, m.nflex, m._impl.nflexedge, m._impl.nflexelem, m._impl.nflexshelldata, m._impl.nflexvert, m.ngeom, m.ngravcomp, m.nhistory, m.njnt, m.nlight, m._impl.nmaxcondim, m._impl.nmaxmeshdeg, m._impl.nmaxpolygon, m._impl.nmaxpyramid, m.nmeshface, m._impl.nrangefinder, m._impl.nsensorcollision, m._impl.nsensorcontact, m._impl.nsensortaxel, m.nsite, m.ntendon, m._impl.ntree, m.nu, m.nv, m._impl.nv_pad, m.nwrap, m._impl.nxn_geom_pair_filtered, m._impl.nxn_pairid, m._impl.nxn_pairid_filtered, m._impl.oct_aabb, m._impl.oct_child, m._impl.oct_coeff, m.pair_dim, m.pair_friction, m.pair_gap, m.pair_margin, m.pair_solimp, m.pair_solref, m.pair_solreffriction, m._impl.plugin, m._impl.plugin_attr, m._impl.qLD_all_updates, m._impl.qLD_level_offsets, m._impl.qLD_updates, m.qpos0, m.qpos_spring, m._impl.rangefinder_sensor_adr, m._impl.sensor_acc_adr, m.sensor_adr, m._impl.sensor_adr_to_contact_adr, m._impl.sensor_contact_adr, m.sensor_cutoff, m.sensor_datatype, m._impl.sensor_delay, m.sensor_dim, m._impl.sensor_e_kinetic, m._impl.sensor_e_potential, m._impl.sensor_history, m._impl.sensor_historyadr, m._impl.sensor_interval, m.sensor_intprm, m._impl.sensor_limitfrc_adr, m._impl.sensor_limitpos_adr, m._impl.sensor_limitvel_adr, m.sensor_objid, m.sensor_objtype, m._impl.sensor_pos_adr, m._impl.sensor_rangefinder_adr, m._impl.sensor_rangefinder_bodyid, m.sensor_refid, m.sensor_reftype, m._impl.sensor_rne_postconstraint, m._impl.sensor_subtree_vel, m._impl.sensor_tendonactfrc_adr, m._impl.sensor_touch_adr, m.sensor_type, m._impl.sensor_vel_adr, m.site_bodyid, m.site_pos, m.site_quat, m.site_size, m.site_type, m._impl.taxel_sensorid, m._impl.taxel_vertadr, m._impl.ten_J_colind, m._impl.ten_J_rowadr, m._impl.ten_J_rownnz, m.tendon_actfrclimited, m.tendon_actfrcrange, m.tendon_adr, m.tendon_armature, m.tendon_damping, m.tendon_dampingpoly, m.tendon_frictionloss, m._impl.tendon_geom_adr, m.tendon_invweight0, m._impl.tendon_jnt_adr, m.tendon_length0, m.tendon_lengthspring, m.tendon_limited, m._impl.tendon_limited_adr, m.tendon_margin, m.tendon_num, m.tendon_range, m._impl.tendon_site_pair_adr, m.tendon_solimp_fri, m.tendon_solimp_lim, m.tendon_solref_fri, m.tendon_solref_lim, m.tendon_stiffness, m.tendon_stiffnesspoly, m._impl.tree_dofadr, m._impl.tree_dofnum, m._impl.tree_sleep_policy, m._impl.wrap_geom_adr, m._impl.wrap_jnt_adr, m.wrap_objid, m.wrap_prm, m._impl.wrap_pulley_scale, m._impl.wrap_site_pair_adr, m.wrap_type, m.opt._impl.broadphase, m.opt._impl.broadphase_filter, m.opt._impl.ccd_iterations, m.opt._impl.ccd_tolerance, m.opt.cone, m.opt._impl.contact_sensor_maxmatch, m.opt.density, m.opt.disableflags, m.opt.enableflags, m.opt._impl.graph_conditional, m.opt.gravity, m.opt._impl.impratio_invsqrt, 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, m.opt._impl.sdf_initpoints, m.opt._impl.sdf_iterations, m.opt.solver, m.opt.timestep, m.opt.tolerance, m.opt.viscosity, m.opt.wind, m.stat.meaninertia, d._impl.naccdmax, d._impl.naconmax, d._impl.njmax, d._impl.njmax_nnz, d._impl.M, d.act, d.act_dot, d.actuator_force, d.actuator_length, d._impl.actuator_moment, d._impl.actuator_velocity, d._impl.body_awake, d._impl.body_awake_ind, d._impl.cacc, d.cam_xmat, d.cam_xpos, d.cdof, d.cdof_dot, d._impl.cfrc_ext, d._impl.cfrc_int, d._impl.cinert, d._impl.crb, d.ctrl, d.cvel, d._impl.dof_awake_ind, d._impl.dof_island, d._impl.dof_islandid, d._impl.efc_islandid, d._impl.energy, d.eq_active, d._impl.flexedge_J, d._impl.flexedge_length, d._impl.flexedge_velocity, d._impl.flexvert_xpos, d.geom_xmat, d.geom_xpos, d.history, d._impl.iqacc, d._impl.iqacc_smooth, d._impl.iqfrc_constraint, d._impl.iqfrc_smooth, d._impl.island_dofadr, d._impl.island_efcadr, d._impl.island_idofadr, d._impl.island_ne, d._impl.island_nefc, d._impl.island_nf, d._impl.island_nv, d._impl.light_xdir, d._impl.light_xpos, d._impl.map_dof2idof, d._impl.map_efc2iefc, d._impl.map_idof2dof, d._impl.map_iefc2efc, d.mocap_pos, d.mocap_quat, d._impl.moment_colind, d._impl.moment_rowadr, d._impl.moment_rownnz, d._impl.nacon, d._impl.nbody_awake, d._impl.ncollision, d._impl.ne, d._impl.nefc, d._impl.nf, d._impl.nidof, d._impl.nisland, d._impl.nl, d._impl.ntree_awake, d._impl.nv_awake, d._impl.qLD, d._impl.qLDiagInv, d.qacc, d.qacc_smooth, d.qacc_warmstart, d.qfrc_actuator, d.qfrc_applied, d.qfrc_bias, d.qfrc_constraint, d._impl.qfrc_damper, d.qfrc_fluid, d.qfrc_gravcomp, d.qfrc_passive, d.qfrc_smooth, d._impl.qfrc_spring, d.qpos, d.qvel, d.sensordata, d.site_xmat, d.site_xpos, d._impl.solver_niter, d._impl.subtree_angmom, d.subtree_com, d._impl.subtree_linvel, d._impl.ten_J, d.ten_length, d._impl.ten_velocity, d._impl.ten_wrapadr, d._impl.ten_wrapnum, d.time, d._impl.tree_asleep, d._impl.tree_awake, d._impl.tree_island, d._impl.wrap_obj, d._impl.wrap_xpos, d.xanchor, d.xaxis, d.xfrc_applied, d.ximat, d.xipos, d.xmat, d.xpos, d.xquat, d._impl.contact__dim, d._impl.contact__dist, d._impl.contact__efc_address, d._impl.contact__flex, d._impl.contact__frame, d._impl.contact__friction, d._impl.contact__geom, d._impl.contact__geomcollisionid, d._impl.contact__includemargin, d._impl.contact__pos, d._impl.contact__solimp, d._impl.contact__solref, d._impl.contact__solreffriction, d._impl.contact__type, d._impl.contact__vert, d._impl.contact__worldid, d._impl.efc__D, d._impl.efc__J, d._impl.efc__J_colind, d._impl.efc__J_rowadr, d._impl.efc__J_rownnz, d._impl.efc__Jqvel, d._impl.efc__Ma, d._impl.efc__aref, d._impl.efc__force, d._impl.efc__frictionloss, d._impl.efc__iD, d._impl.efc__iJ, d._impl.efc__iJ_colind, d._impl.efc__iJ_rowadr, d._impl.efc__iJ_rownnz, d._impl.efc__iaref, d._impl.efc__id, d._impl.efc__iforce, d._impl.efc__ifrictionloss, d._impl.efc__iid, d._impl.efc__island, d._impl.efc__istate, d._impl.efc__itype, d._impl.efc__margin, d._impl.efc__pos, d._impl.efc__state, d._impl.efc__type, d._impl.efc__vel, ) d = d.tree_replace({ '_impl.M': out[0], 'act_dot': out[1], 'actuator_force': out[2], 'actuator_length': out[3], '_impl.actuator_moment': out[4], '_impl.actuator_velocity': out[5], '_impl.body_awake': out[6], '_impl.body_awake_ind': out[7], '_impl.cacc': out[8], 'cam_xmat': out[9], 'cam_xpos': out[10], 'cdof': out[11], 'cdof_dot': out[12], '_impl.cfrc_ext': out[13], '_impl.cfrc_int': out[14], '_impl.cinert': out[15], '_impl.crb': out[16], 'cvel': out[17], '_impl.dof_awake_ind': out[18], '_impl.dof_island': out[19], '_impl.dof_islandid': out[20], '_impl.efc_islandid': out[21], '_impl.energy': out[22], '_impl.flexedge_J': out[23], '_impl.flexedge_length': out[24], '_impl.flexedge_velocity': out[25], '_impl.flexvert_xpos': out[26], 'geom_xmat': out[27], 'geom_xpos': out[28], 'history': out[29], '_impl.iqacc': out[30], '_impl.iqacc_smooth': out[31], '_impl.iqfrc_constraint': out[32], '_impl.iqfrc_smooth': out[33], '_impl.island_dofadr': out[34], '_impl.island_efcadr': out[35], '_impl.island_idofadr': out[36], '_impl.island_ne': out[37], '_impl.island_nefc': out[38], '_impl.island_nf': out[39], '_impl.island_nv': out[40], '_impl.light_xdir': out[41], '_impl.light_xpos': out[42], '_impl.map_dof2idof': out[43], '_impl.map_efc2iefc': out[44], '_impl.map_idof2dof': out[45], '_impl.map_iefc2efc': out[46], '_impl.moment_colind': out[47], '_impl.moment_rowadr': out[48], '_impl.moment_rownnz': out[49], '_impl.nacon': out[50], '_impl.nbody_awake': out[51], '_impl.ncollision': out[52], '_impl.ne': out[53], '_impl.nefc': out[54], '_impl.nf': out[55], '_impl.nidof': out[56], '_impl.nisland': out[57], '_impl.nl': out[58], '_impl.ntree_awake': out[59], '_impl.nv_awake': out[60], '_impl.qLD': out[61], '_impl.qLDiagInv': out[62], 'qacc': out[63], 'qacc_smooth': out[64], 'qfrc_actuator': out[65], 'qfrc_bias': out[66], 'qfrc_constraint': out[67], '_impl.qfrc_damper': out[68], 'qfrc_fluid': out[69], 'qfrc_gravcomp': out[70], 'qfrc_passive': out[71], 'qfrc_smooth': out[72], '_impl.qfrc_spring': out[73], 'qvel': out[74], 'sensordata': out[75], 'site_xmat': out[76], 'site_xpos': out[77], '_impl.solver_niter': out[78], '_impl.subtree_angmom': out[79], 'subtree_com': out[80], '_impl.subtree_linvel': out[81], '_impl.ten_J': out[82], 'ten_length': out[83], '_impl.ten_velocity': out[84], '_impl.ten_wrapadr': out[85], '_impl.ten_wrapnum': out[86], '_impl.tree_asleep': out[87], '_impl.tree_awake': out[88], '_impl.tree_island': out[89], '_impl.wrap_obj': out[90], '_impl.wrap_xpos': out[91], 'xanchor': out[92], 'xaxis': out[93], 'ximat': out[94], 'xipos': out[95], 'xmat': out[96], 'xpos': out[97], 'xquat': out[98], '_impl.contact__dim': out[99], '_impl.contact__dist': out[100], '_impl.contact__efc_address': out[101], '_impl.contact__flex': out[102], '_impl.contact__frame': out[103], '_impl.contact__friction': out[104], '_impl.contact__geom': out[105], '_impl.contact__geomcollisionid': out[106], '_impl.contact__includemargin': out[107], '_impl.contact__pos': out[108], '_impl.contact__solimp': out[109], '_impl.contact__solref': out[110], '_impl.contact__solreffriction': out[111], '_impl.contact__type': out[112], '_impl.contact__vert': out[113], '_impl.contact__worldid': out[114], '_impl.efc__D': out[115], '_impl.efc__J': out[116], '_impl.efc__J_colind': out[117], '_impl.efc__J_rowadr': out[118], '_impl.efc__J_rownnz': out[119], '_impl.efc__Jqvel': out[120], '_impl.efc__Ma': out[121], '_impl.efc__aref': out[122], '_impl.efc__force': out[123], '_impl.efc__frictionloss': out[124], '_impl.efc__iD': out[125], '_impl.efc__iJ': out[126], '_impl.efc__iJ_colind': out[127], '_impl.efc__iJ_rowadr': out[128], '_impl.efc__iJ_rownnz': out[129], '_impl.efc__iaref': out[130], '_impl.efc__id': out[131], '_impl.efc__iforce': out[132], '_impl.efc__ifrictionloss': out[133], '_impl.efc__iid': out[134], '_impl.efc__island': out[135], '_impl.efc__istate': out[136], '_impl.efc__itype': out[137], '_impl.efc__margin': out[138], '_impl.efc__pos': out[139], '_impl.efc__state': out[140], '_impl.efc__type': out[141], '_impl.efc__vel': out[142], }) return d @jax.custom_batching.custom_vmap @ffi.marshal_jax_warp_callable def forward(m: types.Model, d: types.Data): return _forward_jax_impl(m, d) @forward.def_vmap @ffi.marshal_custom_vmap def forward_vmap(unused_axis_size, is_batched, m: types.Model, d: types.Data): d = forward(m, d) return d, is_batched[1] @ffi.format_args_for_warp def _step_shim( # Model nworld: int, D_colind: wp.array[int], D_diag: wp.array[int], D_rowadr: wp.array[int], D_rownnz: wp.array[int], M_elemid: wp.array2d[int], M_fullm_i: wp.array[int], M_fullm_j: wp.array[int], M_mulm_col: wp.array[int], M_mulm_madr: wp.array[int], M_mulm_rowadr: wp.array[int], M_rowadr: wp.array[int], M_rownnz: wp.array[int], M_tiles: tuple[mjwp_types.TileSet, ...], actuator_acc0: wp.array2d[float], actuator_actadr: wp.array[int], actuator_actearly: wp.array[bool], actuator_actlimited: wp.array[bool], actuator_actnum: wp.array[int], actuator_actrange: wp.array2d[wp.vec2], actuator_biasprm: wp.array2d[mjwp_types.vec10f], actuator_biastype: wp.array[int], actuator_cranklength: wp.array2d[float], actuator_ctrllimited: wp.array[bool], actuator_ctrlrange: wp.array2d[wp.vec2], actuator_delay: wp.array[float], actuator_dynprm: wp.array2d[mjwp_types.vec10f], actuator_dyntype: wp.array[int], actuator_forcelimited: wp.array[bool], actuator_forcerange: wp.array2d[wp.vec2], actuator_gainprm: wp.array2d[mjwp_types.vec10f], actuator_gaintype: wp.array[int], actuator_gear: wp.array2d[wp.spatial_vector], actuator_history: wp.array[wp.vec2i], actuator_historyadr: wp.array[int], actuator_lengthrange: wp.array2d[wp.vec2], actuator_trnid: wp.array[wp.vec2i], actuator_trntype: wp.array[int], actuator_trntype_body_adr: wp.array[int], block_dim: mjwp_types.BlockDim, body_branch_start: wp.array[int], body_branches: wp.array[int], body_dofadr: wp.array[int], body_dofnum: wp.array[int], body_fluid_ellipsoid: wp.array[bool], body_geomadr: wp.array[int], body_geomnum: wp.array[int], body_gravcomp: wp.array2d[float], body_inertia: wp.array2d[wp.vec3], body_invweight0: wp.array2d[wp.vec2], body_ipos: wp.array2d[wp.vec3], body_iquat: wp.array2d[wp.quat], body_isdofancestor: wp.array2d[int], body_jntadr: wp.array[int], body_jntnum: wp.array[int], body_mass: wp.array2d[float], body_mocapid: wp.array[int], body_parentid: wp.array[int], body_pos: wp.array2d[wp.vec3], body_quat: wp.array2d[wp.quat], body_rootid: wp.array[int], body_subtreemass: wp.array2d[float], body_tree: tuple[wp.array[int], ...], body_treeid: wp.array[int], body_weldid: wp.array[int], cam_bodyid: wp.array[int], cam_fovy: wp.array2d[float], cam_intrinsic: wp.array2d[wp.vec4], cam_mat0: wp.array2d[wp.mat33], cam_mode: wp.array[int], cam_pos: wp.array2d[wp.vec3], cam_pos0: wp.array2d[wp.vec3], cam_poscom0: wp.array2d[wp.vec3], cam_quat: wp.array2d[wp.quat], cam_resolution: wp.array[wp.vec2i], cam_sensorsize: wp.array[wp.vec2], cam_targetbodyid: wp.array[int], dof_armature: wp.array2d[float], dof_bodyid: wp.array[int], dof_damping: wp.array2d[float], dof_dampingpoly: wp.array2d[wp.vec2], dof_frictionloss: wp.array2d[float], dof_invweight0: wp.array2d[float], dof_jntid: wp.array[int], dof_length: wp.array[float], dof_parentid: wp.array[int], dof_solimp: wp.array2d[mjwp_types.vec5], dof_solref: wp.array2d[wp.vec2], dof_treeid: wp.array[int], dof_tri_col: wp.array[int], dof_tri_row: wp.array[int], eq_connect_adr: wp.array[int], eq_data: wp.array2d[mjwp_types.vec11], eq_flex_adr: wp.array[int], eq_jnt_adr: wp.array[int], eq_obj1id: wp.array[int], eq_obj2id: wp.array[int], eq_objtype: wp.array[int], eq_solimp: wp.array2d[mjwp_types.vec5], eq_solref: wp.array2d[wp.vec2], eq_ten_adr: wp.array[int], eq_type: wp.array[int], eq_wld_adr: wp.array[int], flex_bending: wp.array[float], flex_bendingadr: wp.array[int], flex_centered: wp.array[bool], flex_conaffinity: wp.array[int], flex_condim: wp.array[int], flex_contype: wp.array[int], flex_damping: wp.array[float], flex_dim: wp.array[int], flex_edge: wp.array[wp.vec2i], flex_edgeadr: wp.array[int], flex_edgeflap: wp.array[wp.vec2i], flex_edgenum: wp.array[int], flex_elem: wp.array[int], flex_elemadr: wp.array[int], flex_elemdataadr: wp.array[int], flex_elemedge: wp.array[int], flex_elemedgeadr: wp.array[int], flex_elemnum: wp.array[int], flex_friction: wp.array[wp.vec3], flex_gap: wp.array[float], flex_margin: wp.array[float], flex_priority: wp.array[int], flex_radius: wp.array[float], flex_shell: wp.array[int], flex_shelldataadr: wp.array[int], flex_shellnum: wp.array[int], flex_solimp: wp.array[mjwp_types.vec5], flex_solmix: wp.array[float], flex_solref: wp.array[wp.vec2], flex_stiffness: wp.array[float], flex_stiffnessadr: wp.array[int], flex_vert: wp.array[wp.vec3], flex_vertadr: wp.array[int], flex_vertbodyid: wp.array[int], flex_vertflexid: wp.array[int], flex_vertnum: wp.array[int], flexedge_J_colind: wp.array[int], flexedge_J_rowadr: wp.array[int], flexedge_J_rownnz: wp.array[int], flexedge_invweight0: wp.array[float], flexedge_length0: wp.array[float], geom_aabb: wp.array3d[wp.vec3], geom_bodyid: wp.array[int], geom_conaffinity: wp.array[int], geom_condim: wp.array[int], geom_contype: wp.array[int], geom_dataid: wp.array2d[int], geom_fluid: wp.array2d[float], geom_friction: wp.array2d[wp.vec3], geom_gap: wp.array2d[float], geom_group: wp.array[int], geom_margin: wp.array2d[float], geom_matid: wp.array2d[int], geom_pair_type_count: tuple[int, ...], geom_plugin_index: wp.array[int], geom_pos: wp.array2d[wp.vec3], geom_priority: wp.array[int], geom_quat: wp.array2d[wp.quat], geom_rbound: wp.array2d[float], geom_rgba: wp.array2d[wp.vec4], geom_size: wp.array2d[wp.vec3], geom_solimp: wp.array2d[mjwp_types.vec5], geom_solmix: wp.array2d[float], geom_solref: wp.array2d[wp.vec2], geom_type: wp.array[int], has_fluid: bool, has_sdf_geom: bool, hfield_adr: wp.array[int], hfield_data: wp.array[float], hfield_ncol: wp.array[int], hfield_nrow: wp.array[int], hfield_size: wp.array[wp.vec4], is_sparse: bool, jnt_actfrclimited: wp.array[bool], jnt_actfrcrange: wp.array2d[wp.vec2], jnt_actgravcomp: wp.array[int], jnt_axis: wp.array2d[wp.vec3], jnt_bodyid: wp.array[int], jnt_dofadr: wp.array[int], jnt_limited_ball_adr: wp.array[int], jnt_limited_slide_hinge_adr: wp.array[int], jnt_margin: wp.array2d[float], jnt_pos: wp.array2d[wp.vec3], jnt_qposadr: wp.array[int], jnt_range: wp.array2d[wp.vec2], jnt_solimp: wp.array2d[mjwp_types.vec5], jnt_solref: wp.array2d[wp.vec2], jnt_stiffness: wp.array2d[float], jnt_stiffnesspoly: wp.array2d[wp.vec2], jnt_type: wp.array[int], light_bodyid: wp.array[int], light_dir: wp.array2d[wp.vec3], light_dir0: wp.array2d[wp.vec3], light_mode: wp.array[int], light_pos: wp.array2d[wp.vec3], light_pos0: wp.array2d[wp.vec3], light_poscom0: wp.array2d[wp.vec3], light_targetbodyid: wp.array[int], mapM2D: wp.array[int], mat_rgba: wp.array2d[wp.vec4], max_ten_J_rownnz: int, mesh_face: wp.array[wp.vec3i], mesh_faceadr: wp.array[int], mesh_graph: wp.array[int], mesh_graphadr: wp.array[int], mesh_normal: wp.array[wp.vec3], mesh_normaladr: wp.array[int], mesh_normalnum: wp.array[int], mesh_octadr: wp.array[int], mesh_polyadr: wp.array[int], mesh_polymap: wp.array[int], mesh_polymapadr: wp.array[int], mesh_polymapnum: wp.array[int], mesh_polynormal: wp.array[wp.vec3], mesh_polynum: wp.array[int], mesh_polyvert: wp.array[int], mesh_polyvertadr: wp.array[int], mesh_polyvertnum: wp.array[int], mesh_quat: wp.array[wp.quat], mesh_vert: wp.array[wp.vec3], mesh_vertadr: wp.array[int], mesh_vertnum: wp.array[int], nC: int, nD: int, nJten: int, na: int, nacttrnbody: int, nbody: int, nbranch: int, ncam: int, neq: int, nflex: int, nflexedge: int, nflexelem: int, nflexshelldata: int, nflexvert: int, ngeom: int, ngravcomp: int, nhistory: int, njnt: int, nlight: int, nmaxcondim: int, nmaxmeshdeg: int, nmaxpolygon: int, nmaxpyramid: int, nmeshface: int, nrangefinder: int, nsensorcollision: int, nsensorcontact: int, nsensortaxel: int, nsite: int, ntendon: int, ntree: int, nu: int, nv: int, nv_pad: int, nwrap: int, nxn_geom_pair_filtered: wp.array[wp.vec2i], nxn_pairid: wp.array[wp.vec2i], nxn_pairid_filtered: wp.array[wp.vec2i], oct_aabb: wp.array2d[wp.vec3], oct_child: wp.array[mjwp_types.vec8i], oct_coeff: wp.array[mjwp_types.vec8], pair_dim: wp.array[int], pair_friction: wp.array2d[mjwp_types.vec5], pair_gap: wp.array2d[float], pair_margin: wp.array2d[float], pair_solimp: wp.array2d[mjwp_types.vec5], pair_solref: wp.array2d[wp.vec2], pair_solreffriction: wp.array2d[wp.vec2], plugin: wp.array[int], plugin_attr: wp.array[mjwp_types.vec_pluginattr], qD_fullm_i: wp.array[int], qD_fullm_j: wp.array[int], qLD_all_updates: wp.array[wp.vec3i], qLD_level_offsets: wp.array[int], qLD_updates: tuple[wp.array[wp.vec3i], ...], qpos0: wp.array2d[float], qpos_spring: wp.array2d[float], rangefinder_sensor_adr: wp.array[int], sensor_acc_adr: wp.array[int], sensor_adr: wp.array[int], sensor_adr_to_contact_adr: wp.array[int], sensor_contact_adr: wp.array[int], sensor_cutoff: wp.array[float], sensor_datatype: wp.array[int], sensor_delay: wp.array[float], sensor_dim: wp.array[int], sensor_e_kinetic: bool, sensor_e_potential: bool, sensor_history: wp.array[wp.vec2i], sensor_historyadr: wp.array[int], sensor_interval: wp.array[wp.vec2], sensor_intprm: wp.array2d[int], sensor_limitfrc_adr: wp.array[int], sensor_limitpos_adr: wp.array[int], sensor_limitvel_adr: wp.array[int], sensor_objid: wp.array[int], sensor_objtype: wp.array[int], sensor_pos_adr: wp.array[int], sensor_rangefinder_adr: wp.array[int], sensor_rangefinder_bodyid: wp.array[int], sensor_refid: wp.array[int], sensor_reftype: wp.array[int], sensor_rne_postconstraint: bool, sensor_subtree_vel: bool, sensor_tendonactfrc_adr: wp.array[int], sensor_touch_adr: wp.array[int], sensor_type: wp.array[int], sensor_vel_adr: wp.array[int], site_bodyid: wp.array[int], site_pos: wp.array2d[wp.vec3], site_quat: wp.array2d[wp.quat], site_size: wp.array[wp.vec3], site_type: wp.array[int], taxel_sensorid: wp.array[int], taxel_vertadr: wp.array[int], ten_J_colind: wp.array[int], ten_J_rowadr: wp.array[int], ten_J_rownnz: wp.array[int], tendon_actfrclimited: wp.array[bool], tendon_actfrcrange: wp.array2d[wp.vec2], tendon_adr: wp.array[int], tendon_armature: wp.array2d[float], tendon_damping: wp.array2d[float], tendon_dampingpoly: wp.array2d[wp.vec2], tendon_frictionloss: wp.array2d[float], tendon_geom_adr: wp.array[int], tendon_invweight0: wp.array2d[float], tendon_jnt_adr: wp.array[int], tendon_length0: wp.array2d[float], tendon_lengthspring: wp.array2d[wp.vec2], tendon_limited: wp.array[int], tendon_limited_adr: wp.array[int], tendon_margin: wp.array2d[float], tendon_num: wp.array[int], tendon_range: wp.array2d[wp.vec2], tendon_site_pair_adr: wp.array[int], tendon_solimp_fri: wp.array2d[mjwp_types.vec5], tendon_solimp_lim: wp.array2d[mjwp_types.vec5], tendon_solref_fri: wp.array2d[wp.vec2], tendon_solref_lim: wp.array2d[wp.vec2], tendon_stiffness: wp.array2d[float], tendon_stiffnesspoly: wp.array2d[wp.vec2], tree_dofadr: wp.array[int], tree_dofnum: wp.array[int], tree_sleep_policy: wp.array[int], wrap_geom_adr: wp.array[int], wrap_jnt_adr: wp.array[int], wrap_objid: wp.array[int], wrap_prm: wp.array[float], wrap_pulley_scale: wp.array[float], wrap_site_pair_adr: wp.array[int], wrap_type: wp.array[int], opt__broadphase: int, opt__broadphase_filter: int, opt__ccd_iterations: int, opt__ccd_tolerance: wp.array[float], opt__cone: int, opt__contact_sensor_maxmatch: int, opt__density: wp.array[float], opt__disableflags: int, opt__enableflags: int, opt__graph_conditional: bool, opt__gravity: wp.array[wp.vec3], opt__impratio_invsqrt: wp.array[float], 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, opt__sdf_initpoints: int, opt__sdf_iterations: int, opt__sleep_tolerance: wp.array[float], opt__solver: int, opt__timestep: wp.array[float], opt__tolerance: wp.array[float], opt__viscosity: wp.array[float], opt__wind: wp.array[wp.vec3], stat__meaninertia: wp.array[float], # Data naccdmax: int, naconmax: int, njmax: int, njmax_nnz: int, M: wp.array3d[float], act: wp.array2d[float], act_dot: wp.array2d[float], actuator_force: wp.array2d[float], actuator_length: wp.array2d[float], actuator_moment: wp.array2d[float], actuator_velocity: wp.array2d[float], body_awake: wp.array2d[int], body_awake_ind: wp.array2d[int], cacc: wp.array2d[wp.spatial_vector], cam_xmat: wp.array2d[wp.mat33], cam_xpos: wp.array2d[wp.vec3], cdof: wp.array2d[wp.spatial_vector], cdof_dot: wp.array2d[wp.spatial_vector], cfrc_ext: wp.array2d[wp.spatial_vector], cfrc_int: wp.array2d[wp.spatial_vector], cinert: wp.array2d[mjwp_types.vec10], crb: wp.array2d[mjwp_types.vec10], ctrl: wp.array2d[float], cvel: wp.array2d[wp.spatial_vector], dof_awake_ind: wp.array2d[int], dof_island: wp.array2d[int], dof_islandid: wp.array2d[int], efc_islandid: wp.array2d[int], energy: wp.array[wp.vec2], eq_active: wp.array2d[bool], flexedge_J: wp.array2d[float], flexedge_length: wp.array2d[float], flexedge_velocity: wp.array2d[float], flexvert_xpos: wp.array2d[wp.vec3], geom_xmat: wp.array2d[wp.mat33], geom_xpos: wp.array2d[wp.vec3], history: wp.array2d[float], iqacc: wp.array2d[float], iqacc_smooth: wp.array2d[float], iqfrc_constraint: wp.array2d[float], iqfrc_smooth: wp.array2d[float], island_dofadr: wp.array2d[int], island_efcadr: wp.array2d[int], island_idofadr: wp.array2d[int], island_ne: wp.array2d[int], island_nefc: wp.array2d[int], island_nf: wp.array2d[int], island_nv: wp.array2d[int], light_xdir: wp.array2d[wp.vec3], light_xpos: wp.array2d[wp.vec3], map_dof2idof: wp.array2d[int], map_efc2iefc: wp.array2d[int], map_idof2dof: wp.array2d[int], map_iefc2efc: wp.array2d[int], mocap_pos: wp.array2d[wp.vec3], mocap_quat: wp.array2d[wp.quat], moment_colind: wp.array2d[int], moment_rowadr: wp.array2d[int], moment_rownnz: wp.array2d[int], nacon: wp.array[int], nbody_awake: wp.array[int], ncollision: wp.array[int], ne: wp.array[int], nefc: wp.array[int], nf: wp.array[int], nidof: wp.array[int], nisland: wp.array[int], nl: wp.array[int], ntree_awake: wp.array[int], nv_awake: wp.array[int], qLD: wp.array3d[float], qLDiagInv: wp.array2d[float], qLU: wp.array3d[float], qacc: wp.array2d[float], qacc_smooth: wp.array2d[float], qacc_warmstart: wp.array2d[float], qfrc_actuator: wp.array2d[float], qfrc_applied: wp.array2d[float], qfrc_bias: wp.array2d[float], qfrc_constraint: wp.array2d[float], qfrc_damper: wp.array2d[float], qfrc_fluid: wp.array2d[float], qfrc_gravcomp: wp.array2d[float], qfrc_passive: wp.array2d[float], qfrc_smooth: wp.array2d[float], qfrc_spring: wp.array2d[float], qpos: wp.array2d[float], qvel: wp.array2d[float], sensordata: wp.array2d[float], site_xmat: wp.array2d[wp.mat33], site_xpos: wp.array2d[wp.vec3], solver_niter: wp.array[int], subtree_angmom: wp.array2d[wp.vec3], subtree_com: wp.array2d[wp.vec3], subtree_linvel: wp.array2d[wp.vec3], ten_J: wp.array2d[float], ten_length: wp.array2d[float], ten_velocity: wp.array2d[float], ten_wrapadr: wp.array2d[int], ten_wrapnum: wp.array2d[int], time: wp.array[float], tree_asleep: wp.array2d[int], tree_awake: wp.array2d[int], tree_island: wp.array2d[int], wrap_obj: wp.array2d[wp.vec2i], wrap_xpos: wp.array2d[wp.spatial_vector], xanchor: wp.array2d[wp.vec3], xaxis: wp.array2d[wp.vec3], xfrc_applied: wp.array2d[wp.spatial_vector], ximat: wp.array2d[wp.mat33], xipos: wp.array2d[wp.vec3], xmat: wp.array2d[wp.mat33], xpos: wp.array2d[wp.vec3], xquat: wp.array2d[wp.quat], contact__dim: wp.array[int], contact__dist: wp.array[float], contact__efc_address: wp.array2d[int], contact__flex: wp.array[wp.vec2i], contact__frame: wp.array[wp.mat33], contact__friction: wp.array[mjwp_types.vec5], contact__geom: wp.array[wp.vec2i], contact__geomcollisionid: wp.array[int], contact__includemargin: wp.array[float], contact__pos: wp.array[wp.vec3], contact__solimp: wp.array[mjwp_types.vec5], contact__solref: wp.array[wp.vec2], contact__solreffriction: wp.array[wp.vec2], contact__type: wp.array[int], contact__vert: wp.array[wp.vec2i], contact__worldid: wp.array[int], efc__D: wp.array2d[float], efc__J: wp.array3d[float], efc__J_colind: wp.array3d[int], efc__J_rowadr: wp.array2d[int], efc__J_rownnz: wp.array2d[int], efc__Jqvel: wp.array2d[float], efc__Ma: wp.array2d[float], efc__aref: wp.array2d[float], efc__force: wp.array2d[float], efc__frictionloss: wp.array2d[float], efc__iD: wp.array2d[float], efc__iJ: wp.array3d[float], efc__iJ_colind: wp.array3d[int], efc__iJ_rowadr: wp.array2d[int], efc__iJ_rownnz: wp.array2d[int], efc__iaref: wp.array2d[float], efc__id: wp.array2d[int], efc__iforce: wp.array2d[float], efc__ifrictionloss: wp.array2d[float], efc__iid: wp.array2d[int], efc__island: wp.array2d[int], efc__istate: wp.array2d[int], efc__itype: wp.array2d[int], efc__margin: wp.array2d[float], efc__pos: wp.array2d[float], efc__state: wp.array2d[int], efc__type: wp.array2d[int], efc__vel: wp.array2d[float], ): _m.stat = _s _m.opt = _o _m.callback = _cb _d.efc = _e _d.contact = _c _m.D_colind = D_colind _m.D_diag = D_diag _m.D_rowadr = D_rowadr _m.D_rownnz = D_rownnz _m.M_elemid = M_elemid _m.M_fullm_i = M_fullm_i _m.M_fullm_j = M_fullm_j _m.M_mulm_col = M_mulm_col _m.M_mulm_madr = M_mulm_madr _m.M_mulm_rowadr = M_mulm_rowadr _m.M_rowadr = M_rowadr _m.M_rownnz = M_rownnz _m.M_tiles = M_tiles _m.actuator_acc0 = actuator_acc0 _m.actuator_actadr = actuator_actadr _m.actuator_actearly = actuator_actearly _m.actuator_actlimited = actuator_actlimited _m.actuator_actnum = actuator_actnum _m.actuator_actrange = actuator_actrange _m.actuator_biasprm = actuator_biasprm _m.actuator_biastype = actuator_biastype _m.actuator_cranklength = actuator_cranklength _m.actuator_ctrllimited = actuator_ctrllimited _m.actuator_ctrlrange = actuator_ctrlrange _m.actuator_delay = actuator_delay _m.actuator_dynprm = actuator_dynprm _m.actuator_dyntype = actuator_dyntype _m.actuator_forcelimited = actuator_forcelimited _m.actuator_forcerange = actuator_forcerange _m.actuator_gainprm = actuator_gainprm _m.actuator_gaintype = actuator_gaintype _m.actuator_gear = actuator_gear _m.actuator_history = actuator_history _m.actuator_historyadr = actuator_historyadr _m.actuator_lengthrange = actuator_lengthrange _m.actuator_trnid = actuator_trnid _m.actuator_trntype = actuator_trntype _m.actuator_trntype_body_adr = actuator_trntype_body_adr _m.block_dim = block_dim _m.body_branch_start = body_branch_start _m.body_branches = body_branches _m.body_dofadr = body_dofadr _m.body_dofnum = body_dofnum _m.body_fluid_ellipsoid = body_fluid_ellipsoid _m.body_geomadr = body_geomadr _m.body_geomnum = body_geomnum _m.body_gravcomp = body_gravcomp _m.body_inertia = body_inertia _m.body_invweight0 = body_invweight0 _m.body_ipos = body_ipos _m.body_iquat = body_iquat _m.body_isdofancestor = body_isdofancestor _m.body_jntadr = body_jntadr _m.body_jntnum = body_jntnum _m.body_mass = body_mass _m.body_mocapid = body_mocapid _m.body_parentid = body_parentid _m.body_pos = body_pos _m.body_quat = body_quat _m.body_rootid = body_rootid _m.body_subtreemass = body_subtreemass _m.body_tree = body_tree _m.body_treeid = body_treeid _m.body_weldid = body_weldid _m.cam_bodyid = cam_bodyid _m.cam_fovy = cam_fovy _m.cam_intrinsic = cam_intrinsic _m.cam_mat0 = cam_mat0 _m.cam_mode = cam_mode _m.cam_pos = cam_pos _m.cam_pos0 = cam_pos0 _m.cam_poscom0 = cam_poscom0 _m.cam_quat = cam_quat _m.cam_resolution = cam_resolution _m.cam_sensorsize = cam_sensorsize _m.cam_targetbodyid = cam_targetbodyid _m.dof_armature = dof_armature _m.dof_bodyid = dof_bodyid _m.dof_damping = dof_damping _m.dof_dampingpoly = dof_dampingpoly _m.dof_frictionloss = dof_frictionloss _m.dof_invweight0 = dof_invweight0 _m.dof_jntid = dof_jntid _m.dof_length = dof_length _m.dof_parentid = dof_parentid _m.dof_solimp = dof_solimp _m.dof_solref = dof_solref _m.dof_treeid = dof_treeid _m.dof_tri_col = dof_tri_col _m.dof_tri_row = dof_tri_row _m.eq_connect_adr = eq_connect_adr _m.eq_data = eq_data _m.eq_flex_adr = eq_flex_adr _m.eq_jnt_adr = eq_jnt_adr _m.eq_obj1id = eq_obj1id _m.eq_obj2id = eq_obj2id _m.eq_objtype = eq_objtype _m.eq_solimp = eq_solimp _m.eq_solref = eq_solref _m.eq_ten_adr = eq_ten_adr _m.eq_type = eq_type _m.eq_wld_adr = eq_wld_adr _m.flex_bending = flex_bending _m.flex_bendingadr = flex_bendingadr _m.flex_centered = flex_centered _m.flex_conaffinity = flex_conaffinity _m.flex_condim = flex_condim _m.flex_contype = flex_contype _m.flex_damping = flex_damping _m.flex_dim = flex_dim _m.flex_edge = flex_edge _m.flex_edgeadr = flex_edgeadr _m.flex_edgeflap = flex_edgeflap _m.flex_edgenum = flex_edgenum _m.flex_elem = flex_elem _m.flex_elemadr = flex_elemadr _m.flex_elemdataadr = flex_elemdataadr _m.flex_elemedge = flex_elemedge _m.flex_elemedgeadr = flex_elemedgeadr _m.flex_elemnum = flex_elemnum _m.flex_friction = flex_friction _m.flex_gap = flex_gap _m.flex_margin = flex_margin _m.flex_priority = flex_priority _m.flex_radius = flex_radius _m.flex_shell = flex_shell _m.flex_shelldataadr = flex_shelldataadr _m.flex_shellnum = flex_shellnum _m.flex_solimp = flex_solimp _m.flex_solmix = flex_solmix _m.flex_solref = flex_solref _m.flex_stiffness = flex_stiffness _m.flex_stiffnessadr = flex_stiffnessadr _m.flex_vert = flex_vert _m.flex_vertadr = flex_vertadr _m.flex_vertbodyid = flex_vertbodyid _m.flex_vertflexid = flex_vertflexid _m.flex_vertnum = flex_vertnum _m.flexedge_J_colind = flexedge_J_colind _m.flexedge_J_rowadr = flexedge_J_rowadr _m.flexedge_J_rownnz = flexedge_J_rownnz _m.flexedge_invweight0 = flexedge_invweight0 _m.flexedge_length0 = flexedge_length0 _m.geom_aabb = geom_aabb _m.geom_bodyid = geom_bodyid _m.geom_conaffinity = geom_conaffinity _m.geom_condim = geom_condim _m.geom_contype = geom_contype _m.geom_dataid = geom_dataid _m.geom_fluid = geom_fluid _m.geom_friction = geom_friction _m.geom_gap = geom_gap _m.geom_group = geom_group _m.geom_margin = geom_margin _m.geom_matid = geom_matid _m.geom_pair_type_count = geom_pair_type_count _m.geom_plugin_index = geom_plugin_index _m.geom_pos = geom_pos _m.geom_priority = geom_priority _m.geom_quat = geom_quat _m.geom_rbound = geom_rbound _m.geom_rgba = geom_rgba _m.geom_size = geom_size _m.geom_solimp = geom_solimp _m.geom_solmix = geom_solmix _m.geom_solref = geom_solref _m.geom_type = geom_type _m.has_fluid = has_fluid _m.has_sdf_geom = has_sdf_geom _m.hfield_adr = hfield_adr _m.hfield_data = hfield_data _m.hfield_ncol = hfield_ncol _m.hfield_nrow = hfield_nrow _m.hfield_size = hfield_size _m.is_sparse = is_sparse _m.jnt_actfrclimited = jnt_actfrclimited _m.jnt_actfrcrange = jnt_actfrcrange _m.jnt_actgravcomp = jnt_actgravcomp _m.jnt_axis = jnt_axis _m.jnt_bodyid = jnt_bodyid _m.jnt_dofadr = jnt_dofadr _m.jnt_limited_ball_adr = jnt_limited_ball_adr _m.jnt_limited_slide_hinge_adr = jnt_limited_slide_hinge_adr _m.jnt_margin = jnt_margin _m.jnt_pos = jnt_pos _m.jnt_qposadr = jnt_qposadr _m.jnt_range = jnt_range _m.jnt_solimp = jnt_solimp _m.jnt_solref = jnt_solref _m.jnt_stiffness = jnt_stiffness _m.jnt_stiffnesspoly = jnt_stiffnesspoly _m.jnt_type = jnt_type _m.light_bodyid = light_bodyid _m.light_dir = light_dir _m.light_dir0 = light_dir0 _m.light_mode = light_mode _m.light_pos = light_pos _m.light_pos0 = light_pos0 _m.light_poscom0 = light_poscom0 _m.light_targetbodyid = light_targetbodyid _m.mapM2D = mapM2D _m.mat_rgba = mat_rgba _m.max_ten_J_rownnz = max_ten_J_rownnz _m.mesh_face = mesh_face _m.mesh_faceadr = mesh_faceadr _m.mesh_graph = mesh_graph _m.mesh_graphadr = mesh_graphadr _m.mesh_normal = mesh_normal _m.mesh_normaladr = mesh_normaladr _m.mesh_normalnum = mesh_normalnum _m.mesh_octadr = mesh_octadr _m.mesh_polyadr = mesh_polyadr _m.mesh_polymap = mesh_polymap _m.mesh_polymapadr = mesh_polymapadr _m.mesh_polymapnum = mesh_polymapnum _m.mesh_polynormal = mesh_polynormal _m.mesh_polynum = mesh_polynum _m.mesh_polyvert = mesh_polyvert _m.mesh_polyvertadr = mesh_polyvertadr _m.mesh_polyvertnum = mesh_polyvertnum _m.mesh_quat = mesh_quat _m.mesh_vert = mesh_vert _m.mesh_vertadr = mesh_vertadr _m.mesh_vertnum = mesh_vertnum _m.nC = nC _m.nD = nD _m.nJten = nJten _m.na = na _m.nacttrnbody = nacttrnbody _m.nbody = nbody _m.nbranch = nbranch _m.ncam = ncam _m.neq = neq _m.nflex = nflex _m.nflexedge = nflexedge _m.nflexelem = nflexelem _m.nflexshelldata = nflexshelldata _m.nflexvert = nflexvert _m.ngeom = ngeom _m.ngravcomp = ngravcomp _m.nhistory = nhistory _m.njnt = njnt _m.nlight = nlight _m.nmaxcondim = nmaxcondim _m.nmaxmeshdeg = nmaxmeshdeg _m.nmaxpolygon = nmaxpolygon _m.nmaxpyramid = nmaxpyramid _m.nmeshface = nmeshface _m.nrangefinder = nrangefinder _m.nsensorcollision = nsensorcollision _m.nsensorcontact = nsensorcontact _m.nsensortaxel = nsensortaxel _m.nsite = nsite _m.ntendon = ntendon _m.ntree = ntree _m.nu = nu _m.nv = nv _m.nv_pad = nv_pad _m.nwrap = nwrap _m.nxn_geom_pair_filtered = nxn_geom_pair_filtered _m.nxn_pairid = nxn_pairid _m.nxn_pairid_filtered = nxn_pairid_filtered _m.oct_aabb = oct_aabb _m.oct_child = oct_child _m.oct_coeff = oct_coeff _m.opt.broadphase = opt__broadphase _m.opt.broadphase_filter = opt__broadphase_filter _m.opt.ccd_iterations = opt__ccd_iterations _m.opt.ccd_tolerance = opt__ccd_tolerance _m.opt.cone = opt__cone _m.opt.contact_sensor_maxmatch = opt__contact_sensor_maxmatch _m.opt.density = opt__density _m.opt.disableflags = opt__disableflags _m.opt.enableflags = opt__enableflags _m.opt.graph_conditional = opt__graph_conditional _m.opt.gravity = opt__gravity _m.opt.impratio_invsqrt = opt__impratio_invsqrt _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 _m.opt.sdf_initpoints = opt__sdf_initpoints _m.opt.sdf_iterations = opt__sdf_iterations _m.opt.sleep_tolerance = opt__sleep_tolerance _m.opt.solver = opt__solver _m.opt.timestep = opt__timestep _m.opt.tolerance = opt__tolerance _m.opt.viscosity = opt__viscosity _m.opt.wind = opt__wind _m.pair_dim = pair_dim _m.pair_friction = pair_friction _m.pair_gap = pair_gap _m.pair_margin = pair_margin _m.pair_solimp = pair_solimp _m.pair_solref = pair_solref _m.pair_solreffriction = pair_solreffriction _m.plugin = plugin _m.plugin_attr = plugin_attr _m.qD_fullm_i = qD_fullm_i _m.qD_fullm_j = qD_fullm_j _m.qLD_all_updates = qLD_all_updates _m.qLD_level_offsets = qLD_level_offsets _m.qLD_updates = qLD_updates _m.qpos0 = qpos0 _m.qpos_spring = qpos_spring _m.rangefinder_sensor_adr = rangefinder_sensor_adr _m.sensor_acc_adr = sensor_acc_adr _m.sensor_adr = sensor_adr _m.sensor_adr_to_contact_adr = sensor_adr_to_contact_adr _m.sensor_contact_adr = sensor_contact_adr _m.sensor_cutoff = sensor_cutoff _m.sensor_datatype = sensor_datatype _m.sensor_delay = sensor_delay _m.sensor_dim = sensor_dim _m.sensor_e_kinetic = sensor_e_kinetic _m.sensor_e_potential = sensor_e_potential _m.sensor_history = sensor_history _m.sensor_historyadr = sensor_historyadr _m.sensor_interval = sensor_interval _m.sensor_intprm = sensor_intprm _m.sensor_limitfrc_adr = sensor_limitfrc_adr _m.sensor_limitpos_adr = sensor_limitpos_adr _m.sensor_limitvel_adr = sensor_limitvel_adr _m.sensor_objid = sensor_objid _m.sensor_objtype = sensor_objtype _m.sensor_pos_adr = sensor_pos_adr _m.sensor_rangefinder_adr = sensor_rangefinder_adr _m.sensor_rangefinder_bodyid = sensor_rangefinder_bodyid _m.sensor_refid = sensor_refid _m.sensor_reftype = sensor_reftype _m.sensor_rne_postconstraint = sensor_rne_postconstraint _m.sensor_subtree_vel = sensor_subtree_vel _m.sensor_tendonactfrc_adr = sensor_tendonactfrc_adr _m.sensor_touch_adr = sensor_touch_adr _m.sensor_type = sensor_type _m.sensor_vel_adr = sensor_vel_adr _m.site_bodyid = site_bodyid _m.site_pos = site_pos _m.site_quat = site_quat _m.site_size = site_size _m.site_type = site_type _m.stat.meaninertia = stat__meaninertia _m.taxel_sensorid = taxel_sensorid _m.taxel_vertadr = taxel_vertadr _m.ten_J_colind = ten_J_colind _m.ten_J_rowadr = ten_J_rowadr _m.ten_J_rownnz = ten_J_rownnz _m.tendon_actfrclimited = tendon_actfrclimited _m.tendon_actfrcrange = tendon_actfrcrange _m.tendon_adr = tendon_adr _m.tendon_armature = tendon_armature _m.tendon_damping = tendon_damping _m.tendon_dampingpoly = tendon_dampingpoly _m.tendon_frictionloss = tendon_frictionloss _m.tendon_geom_adr = tendon_geom_adr _m.tendon_invweight0 = tendon_invweight0 _m.tendon_jnt_adr = tendon_jnt_adr _m.tendon_length0 = tendon_length0 _m.tendon_lengthspring = tendon_lengthspring _m.tendon_limited = tendon_limited _m.tendon_limited_adr = tendon_limited_adr _m.tendon_margin = tendon_margin _m.tendon_num = tendon_num _m.tendon_range = tendon_range _m.tendon_site_pair_adr = tendon_site_pair_adr _m.tendon_solimp_fri = tendon_solimp_fri _m.tendon_solimp_lim = tendon_solimp_lim _m.tendon_solref_fri = tendon_solref_fri _m.tendon_solref_lim = tendon_solref_lim _m.tendon_stiffness = tendon_stiffness _m.tendon_stiffnesspoly = tendon_stiffnesspoly _m.tree_dofadr = tree_dofadr _m.tree_dofnum = tree_dofnum _m.tree_sleep_policy = tree_sleep_policy _m.wrap_geom_adr = wrap_geom_adr _m.wrap_jnt_adr = wrap_jnt_adr _m.wrap_objid = wrap_objid _m.wrap_prm = wrap_prm _m.wrap_pulley_scale = wrap_pulley_scale _m.wrap_site_pair_adr = wrap_site_pair_adr _m.wrap_type = wrap_type _d.M = M _d.act = act _d.act_dot = act_dot _d.actuator_force = actuator_force _d.actuator_length = actuator_length _d.actuator_moment = actuator_moment _d.actuator_velocity = actuator_velocity _d.body_awake = body_awake _d.body_awake_ind = body_awake_ind _d.cacc = cacc _d.cam_xmat = cam_xmat _d.cam_xpos = cam_xpos _d.cdof = cdof _d.cdof_dot = cdof_dot _d.cfrc_ext = cfrc_ext _d.cfrc_int = cfrc_int _d.cinert = cinert _d.contact.dim = contact__dim _d.contact.dist = contact__dist _d.contact.efc_address = contact__efc_address _d.contact.flex = contact__flex _d.contact.frame = contact__frame _d.contact.friction = contact__friction _d.contact.geom = contact__geom _d.contact.geomcollisionid = contact__geomcollisionid _d.contact.includemargin = contact__includemargin _d.contact.pos = contact__pos _d.contact.solimp = contact__solimp _d.contact.solref = contact__solref _d.contact.solreffriction = contact__solreffriction _d.contact.type = contact__type _d.contact.vert = contact__vert _d.contact.worldid = contact__worldid _d.crb = crb _d.ctrl = ctrl _d.cvel = cvel _d.dof_awake_ind = dof_awake_ind _d.dof_island = dof_island _d.dof_islandid = dof_islandid _d.efc.D = efc__D _d.efc.J = efc__J _d.efc.J_colind = efc__J_colind _d.efc.J_rowadr = efc__J_rowadr _d.efc.J_rownnz = efc__J_rownnz _d.efc.Jqvel = efc__Jqvel _d.efc.Ma = efc__Ma _d.efc.aref = efc__aref _d.efc.force = efc__force _d.efc.frictionloss = efc__frictionloss _d.efc.iD = efc__iD _d.efc.iJ = efc__iJ _d.efc.iJ_colind = efc__iJ_colind _d.efc.iJ_rowadr = efc__iJ_rowadr _d.efc.iJ_rownnz = efc__iJ_rownnz _d.efc.iaref = efc__iaref _d.efc.id = efc__id _d.efc.iforce = efc__iforce _d.efc.ifrictionloss = efc__ifrictionloss _d.efc.iid = efc__iid _d.efc.island = efc__island _d.efc.istate = efc__istate _d.efc.itype = efc__itype _d.efc.margin = efc__margin _d.efc.pos = efc__pos _d.efc.state = efc__state _d.efc.type = efc__type _d.efc.vel = efc__vel _d.efc_islandid = efc_islandid _d.energy = energy _d.eq_active = eq_active _d.flexedge_J = flexedge_J _d.flexedge_length = flexedge_length _d.flexedge_velocity = flexedge_velocity _d.flexvert_xpos = flexvert_xpos _d.geom_xmat = geom_xmat _d.geom_xpos = geom_xpos _d.history = history _d.iqacc = iqacc _d.iqacc_smooth = iqacc_smooth _d.iqfrc_constraint = iqfrc_constraint _d.iqfrc_smooth = iqfrc_smooth _d.island_dofadr = island_dofadr _d.island_efcadr = island_efcadr _d.island_idofadr = island_idofadr _d.island_ne = island_ne _d.island_nefc = island_nefc _d.island_nf = island_nf _d.island_nv = island_nv _d.light_xdir = light_xdir _d.light_xpos = light_xpos _d.map_dof2idof = map_dof2idof _d.map_efc2iefc = map_efc2iefc _d.map_idof2dof = map_idof2dof _d.map_iefc2efc = map_iefc2efc _d.mocap_pos = mocap_pos _d.mocap_quat = mocap_quat _d.moment_colind = moment_colind _d.moment_rowadr = moment_rowadr _d.moment_rownnz = moment_rownnz _d.naccdmax = naccdmax _d.nacon = nacon _d.naconmax = naconmax _d.nbody_awake = nbody_awake _d.ncollision = ncollision _d.ne = ne _d.nefc = nefc _d.nf = nf _d.nidof = nidof _d.nisland = nisland _d.njmax = njmax _d.njmax_nnz = njmax_nnz _d.nl = nl _d.ntree_awake = ntree_awake _d.nv_awake = nv_awake _d.qLD = qLD _d.qLDiagInv = qLDiagInv _d.qLU = qLU _d.qacc = qacc _d.qacc_smooth = qacc_smooth _d.qacc_warmstart = qacc_warmstart _d.qfrc_actuator = qfrc_actuator _d.qfrc_applied = qfrc_applied _d.qfrc_bias = qfrc_bias _d.qfrc_constraint = qfrc_constraint _d.qfrc_damper = qfrc_damper _d.qfrc_fluid = qfrc_fluid _d.qfrc_gravcomp = qfrc_gravcomp _d.qfrc_passive = qfrc_passive _d.qfrc_smooth = qfrc_smooth _d.qfrc_spring = qfrc_spring _d.qpos = qpos _d.qvel = qvel _d.sensordata = sensordata _d.site_xmat = site_xmat _d.site_xpos = site_xpos _d.solver_niter = solver_niter _d.subtree_angmom = subtree_angmom _d.subtree_com = subtree_com _d.subtree_linvel = subtree_linvel _d.ten_J = ten_J _d.ten_length = ten_length _d.ten_velocity = ten_velocity _d.ten_wrapadr = ten_wrapadr _d.ten_wrapnum = ten_wrapnum _d.time = time _d.tree_asleep = tree_asleep _d.tree_awake = tree_awake _d.tree_island = tree_island _d.wrap_obj = wrap_obj _d.wrap_xpos = wrap_xpos _d.xanchor = xanchor _d.xaxis = xaxis _d.xfrc_applied = xfrc_applied _d.ximat = ximat _d.xipos = xipos _d.xmat = xmat _d.xpos = xpos _d.xquat = xquat _d.nworld = nworld mjwarp.step(_m, _d) def _step_jax_impl(m: types.Model, d: types.Data): output_dims = { 'M': d._impl.M.shape, 'act': d.act.shape, 'act_dot': d.act_dot.shape, 'actuator_force': d.actuator_force.shape, 'actuator_length': d.actuator_length.shape, 'actuator_moment': d._impl.actuator_moment.shape, 'actuator_velocity': d._impl.actuator_velocity.shape, 'body_awake': d._impl.body_awake.shape, 'body_awake_ind': d._impl.body_awake_ind.shape, 'cacc': d._impl.cacc.shape, 'cam_xmat': d.cam_xmat.shape, 'cam_xpos': d.cam_xpos.shape, 'cdof': d.cdof.shape, 'cdof_dot': d.cdof_dot.shape, 'cfrc_ext': d._impl.cfrc_ext.shape, 'cfrc_int': d._impl.cfrc_int.shape, 'cinert': d._impl.cinert.shape, 'crb': d._impl.crb.shape, 'cvel': d.cvel.shape, 'dof_awake_ind': d._impl.dof_awake_ind.shape, 'dof_island': d._impl.dof_island.shape, 'dof_islandid': d._impl.dof_islandid.shape, 'efc_islandid': d._impl.efc_islandid.shape, 'energy': d._impl.energy.shape, 'flexedge_J': d._impl.flexedge_J.shape, 'flexedge_length': d._impl.flexedge_length.shape, 'flexedge_velocity': d._impl.flexedge_velocity.shape, 'flexvert_xpos': d._impl.flexvert_xpos.shape, 'geom_xmat': d.geom_xmat.shape, 'geom_xpos': d.geom_xpos.shape, 'history': d.history.shape, 'iqacc': d._impl.iqacc.shape, 'iqacc_smooth': d._impl.iqacc_smooth.shape, 'iqfrc_constraint': d._impl.iqfrc_constraint.shape, 'iqfrc_smooth': d._impl.iqfrc_smooth.shape, 'island_dofadr': d._impl.island_dofadr.shape, 'island_efcadr': d._impl.island_efcadr.shape, 'island_idofadr': d._impl.island_idofadr.shape, 'island_ne': d._impl.island_ne.shape, 'island_nefc': d._impl.island_nefc.shape, 'island_nf': d._impl.island_nf.shape, 'island_nv': d._impl.island_nv.shape, 'light_xdir': d._impl.light_xdir.shape, 'light_xpos': d._impl.light_xpos.shape, 'map_dof2idof': d._impl.map_dof2idof.shape, 'map_efc2iefc': d._impl.map_efc2iefc.shape, 'map_idof2dof': d._impl.map_idof2dof.shape, 'map_iefc2efc': d._impl.map_iefc2efc.shape, 'moment_colind': d._impl.moment_colind.shape, 'moment_rowadr': d._impl.moment_rowadr.shape, 'moment_rownnz': d._impl.moment_rownnz.shape, 'nacon': d._impl.nacon.shape, 'nbody_awake': d._impl.nbody_awake.shape, 'ncollision': d._impl.ncollision.shape, 'ne': d._impl.ne.shape, 'nefc': d._impl.nefc.shape, 'nf': d._impl.nf.shape, 'nidof': d._impl.nidof.shape, 'nisland': d._impl.nisland.shape, 'nl': d._impl.nl.shape, 'ntree_awake': d._impl.ntree_awake.shape, 'nv_awake': d._impl.nv_awake.shape, 'qLD': d._impl.qLD.shape, 'qLDiagInv': d._impl.qLDiagInv.shape, 'qLU': d._impl.qLU.shape, 'qacc': d.qacc.shape, 'qacc_smooth': d.qacc_smooth.shape, 'qacc_warmstart': d.qacc_warmstart.shape, 'qfrc_actuator': d.qfrc_actuator.shape, 'qfrc_bias': d.qfrc_bias.shape, 'qfrc_constraint': d.qfrc_constraint.shape, 'qfrc_damper': d._impl.qfrc_damper.shape, 'qfrc_fluid': d.qfrc_fluid.shape, 'qfrc_gravcomp': d.qfrc_gravcomp.shape, 'qfrc_passive': d.qfrc_passive.shape, 'qfrc_smooth': d.qfrc_smooth.shape, 'qfrc_spring': d._impl.qfrc_spring.shape, 'qpos': d.qpos.shape, 'qvel': d.qvel.shape, 'sensordata': d.sensordata.shape, 'site_xmat': d.site_xmat.shape, 'site_xpos': d.site_xpos.shape, 'solver_niter': d._impl.solver_niter.shape, 'subtree_angmom': d._impl.subtree_angmom.shape, 'subtree_com': d.subtree_com.shape, 'subtree_linvel': d._impl.subtree_linvel.shape, 'ten_J': d._impl.ten_J.shape, 'ten_length': d.ten_length.shape, 'ten_velocity': d._impl.ten_velocity.shape, 'ten_wrapadr': d._impl.ten_wrapadr.shape, 'ten_wrapnum': d._impl.ten_wrapnum.shape, 'time': d.time.shape, 'tree_asleep': d._impl.tree_asleep.shape, 'tree_awake': d._impl.tree_awake.shape, 'tree_island': d._impl.tree_island.shape, 'wrap_obj': d._impl.wrap_obj.shape, 'wrap_xpos': d._impl.wrap_xpos.shape, 'xanchor': d.xanchor.shape, 'xaxis': d.xaxis.shape, 'ximat': d.ximat.shape, 'xipos': d.xipos.shape, 'xmat': d.xmat.shape, 'xpos': d.xpos.shape, 'xquat': d.xquat.shape, 'contact__dim': d._impl.contact__dim.shape, 'contact__dist': d._impl.contact__dist.shape, 'contact__efc_address': d._impl.contact__efc_address.shape, 'contact__flex': d._impl.contact__flex.shape, 'contact__frame': d._impl.contact__frame.shape, 'contact__friction': d._impl.contact__friction.shape, 'contact__geom': d._impl.contact__geom.shape, 'contact__geomcollisionid': d._impl.contact__geomcollisionid.shape, 'contact__includemargin': d._impl.contact__includemargin.shape, 'contact__pos': d._impl.contact__pos.shape, 'contact__solimp': d._impl.contact__solimp.shape, 'contact__solref': d._impl.contact__solref.shape, 'contact__solreffriction': d._impl.contact__solreffriction.shape, 'contact__type': d._impl.contact__type.shape, 'contact__vert': d._impl.contact__vert.shape, 'contact__worldid': d._impl.contact__worldid.shape, 'efc__D': d._impl.efc__D.shape, 'efc__J': d._impl.efc__J.shape, 'efc__J_colind': d._impl.efc__J_colind.shape, 'efc__J_rowadr': d._impl.efc__J_rowadr.shape, 'efc__J_rownnz': d._impl.efc__J_rownnz.shape, 'efc__Jqvel': d._impl.efc__Jqvel.shape, 'efc__Ma': d._impl.efc__Ma.shape, 'efc__aref': d._impl.efc__aref.shape, 'efc__force': d._impl.efc__force.shape, 'efc__frictionloss': d._impl.efc__frictionloss.shape, 'efc__iD': d._impl.efc__iD.shape, 'efc__iJ': d._impl.efc__iJ.shape, 'efc__iJ_colind': d._impl.efc__iJ_colind.shape, 'efc__iJ_rowadr': d._impl.efc__iJ_rowadr.shape, 'efc__iJ_rownnz': d._impl.efc__iJ_rownnz.shape, 'efc__iaref': d._impl.efc__iaref.shape, 'efc__id': d._impl.efc__id.shape, 'efc__iforce': d._impl.efc__iforce.shape, 'efc__ifrictionloss': d._impl.efc__ifrictionloss.shape, 'efc__iid': d._impl.efc__iid.shape, 'efc__island': d._impl.efc__island.shape, 'efc__istate': d._impl.efc__istate.shape, 'efc__itype': d._impl.efc__itype.shape, 'efc__margin': d._impl.efc__margin.shape, 'efc__pos': d._impl.efc__pos.shape, 'efc__state': d._impl.efc__state.shape, 'efc__type': d._impl.efc__type.shape, 'efc__vel': d._impl.efc__vel.shape, } jf = ffi.jax_callable_variadic_tuple( _step_shim, num_outputs=148, output_dims=output_dims, vmap_method=None, in_out_argnames=set([ 'M', 'act', 'act_dot', 'actuator_force', 'actuator_length', 'actuator_moment', 'actuator_velocity', 'body_awake', 'body_awake_ind', 'cacc', 'cam_xmat', 'cam_xpos', 'cdof', 'cdof_dot', 'cfrc_ext', 'cfrc_int', 'cinert', 'crb', 'cvel', 'dof_awake_ind', 'dof_island', 'dof_islandid', 'efc_islandid', 'energy', 'flexedge_J', 'flexedge_length', 'flexedge_velocity', 'flexvert_xpos', 'geom_xmat', 'geom_xpos', 'history', 'iqacc', 'iqacc_smooth', 'iqfrc_constraint', 'iqfrc_smooth', 'island_dofadr', 'island_efcadr', 'island_idofadr', 'island_ne', 'island_nefc', 'island_nf', 'island_nv', 'light_xdir', 'light_xpos', 'map_dof2idof', 'map_efc2iefc', 'map_idof2dof', 'map_iefc2efc', 'moment_colind', 'moment_rowadr', 'moment_rownnz', 'nacon', 'nbody_awake', 'ncollision', 'ne', 'nefc', 'nf', 'nidof', 'nisland', 'nl', 'ntree_awake', 'nv_awake', 'qLD', 'qLDiagInv', 'qLU', 'qacc', 'qacc_smooth', 'qacc_warmstart', 'qfrc_actuator', 'qfrc_bias', 'qfrc_constraint', 'qfrc_damper', 'qfrc_fluid', 'qfrc_gravcomp', 'qfrc_passive', 'qfrc_smooth', 'qfrc_spring', 'qpos', 'qvel', 'sensordata', 'site_xmat', 'site_xpos', 'solver_niter', 'subtree_angmom', 'subtree_com', 'subtree_linvel', 'ten_J', 'ten_length', 'ten_velocity', 'ten_wrapadr', 'ten_wrapnum', 'time', 'tree_asleep', 'tree_awake', 'tree_island', 'wrap_obj', 'wrap_xpos', 'xanchor', 'xaxis', 'ximat', 'xipos', 'xmat', 'xpos', 'xquat', 'contact__dim', 'contact__dist', 'contact__efc_address', 'contact__flex', 'contact__frame', 'contact__friction', 'contact__geom', 'contact__geomcollisionid', 'contact__includemargin', 'contact__pos', 'contact__solimp', 'contact__solref', 'contact__solreffriction', 'contact__type', 'contact__vert', 'contact__worldid', 'efc__D', 'efc__J', 'efc__J_colind', 'efc__J_rowadr', 'efc__J_rownnz', 'efc__Jqvel', 'efc__Ma', 'efc__aref', 'efc__force', 'efc__frictionloss', 'efc__iD', 'efc__iJ', 'efc__iJ_colind', 'efc__iJ_rowadr', 'efc__iJ_rownnz', 'efc__iaref', 'efc__id', 'efc__iforce', 'efc__ifrictionloss', 'efc__iid', 'efc__island', 'efc__istate', 'efc__itype', 'efc__margin', 'efc__pos', 'efc__state', 'efc__type', 'efc__vel', ]), stage_in_argnames=set([ 'act', 'act_dot', 'actuator_acc0', 'actuator_actrange', 'actuator_biasprm', 'actuator_cranklength', 'actuator_ctrlrange', 'actuator_dynprm', 'actuator_force', 'actuator_forcerange', 'actuator_gainprm', 'actuator_gear', 'actuator_length', 'actuator_lengthrange', 'body_gravcomp', 'body_inertia', 'body_invweight0', 'body_ipos', 'body_iquat', 'body_mass', 'body_pos', 'body_quat', 'body_subtreemass', 'cam_fovy', 'cam_intrinsic', 'cam_mat0', 'cam_pos', 'cam_pos0', 'cam_poscom0', 'cam_quat', 'cam_xmat', 'cam_xpos', 'cdof', 'cdof_dot', 'ctrl', 'cvel', 'dof_armature', 'dof_damping', 'dof_dampingpoly', 'dof_frictionloss', 'dof_invweight0', 'dof_solimp', 'dof_solref', 'eq_active', 'eq_data', 'eq_solimp', 'eq_solref', 'geom_aabb', 'geom_friction', 'geom_gap', 'geom_margin', 'geom_matid', 'geom_pos', 'geom_quat', 'geom_rbound', 'geom_rgba', 'geom_size', 'geom_solimp', 'geom_solmix', 'geom_solref', 'geom_xmat', 'geom_xpos', 'hfield_data', 'history', 'jnt_actfrcrange', 'jnt_axis', 'jnt_margin', 'jnt_pos', 'jnt_range', 'jnt_solimp', 'jnt_solref', 'jnt_stiffness', 'jnt_stiffnesspoly', 'light_dir', 'light_dir0', 'light_pos', 'light_pos0', 'light_poscom0', 'mat_rgba', 'mocap_pos', 'mocap_quat', 'opt__density', 'opt__gravity', 'opt__ls_tolerance', 'opt__magnetic', 'opt__timestep', 'opt__tolerance', 'opt__viscosity', 'opt__wind', 'pair_friction', 'pair_gap', 'pair_margin', 'pair_solimp', 'pair_solref', 'pair_solreffriction', 'qacc', 'qacc_smooth', 'qacc_warmstart', 'qfrc_actuator', 'qfrc_applied', 'qfrc_bias', 'qfrc_constraint', 'qfrc_fluid', 'qfrc_gravcomp', 'qfrc_passive', 'qfrc_smooth', 'qpos', 'qpos0', 'qpos_spring', 'qvel', 'sensordata', 'site_pos', 'site_quat', 'site_xmat', 'site_xpos', 'subtree_com', 'ten_length', 'tendon_actfrcrange', 'tendon_armature', 'tendon_damping', 'tendon_dampingpoly', 'tendon_frictionloss', 'tendon_invweight0', 'tendon_length0', 'tendon_lengthspring', 'tendon_margin', 'tendon_range', 'tendon_solimp_fri', 'tendon_solimp_lim', 'tendon_solref_fri', 'tendon_solref_lim', 'tendon_stiffness', 'tendon_stiffnesspoly', 'time', 'xanchor', 'xaxis', 'xfrc_applied', 'ximat', 'xipos', 'xmat', 'xpos', 'xquat', ]), stage_out_argnames=set([ 'act', 'act_dot', 'actuator_force', 'actuator_length', 'cam_xmat', 'cam_xpos', 'cdof', 'cdof_dot', 'cvel', 'geom_xmat', 'geom_xpos', 'history', 'qacc', 'qacc_smooth', 'qacc_warmstart', 'qfrc_actuator', 'qfrc_bias', 'qfrc_constraint', 'qfrc_fluid', 'qfrc_gravcomp', 'qfrc_passive', 'qfrc_smooth', 'qpos', 'qvel', 'sensordata', 'site_xmat', 'site_xpos', 'subtree_com', 'ten_length', 'time', 'xanchor', 'xaxis', 'ximat', 'xipos', 'xmat', 'xpos', 'xquat', ]), graph_mode=m.opt._impl.graph_mode, has_side_effect=False, ) out = jf( d.qpos.shape[0], m._impl.D_colind, m._impl.D_diag, m._impl.D_rowadr, m._impl.D_rownnz, m._impl.M_elemid, m._impl.M_fullm_i, m._impl.M_fullm_j, m._impl.M_mulm_col, m._impl.M_mulm_madr, m._impl.M_mulm_rowadr, m._impl.M_rowadr, m._impl.M_rownnz, m._impl.M_tiles, m.actuator_acc0, m.actuator_actadr, m.actuator_actearly, m.actuator_actlimited, m.actuator_actnum, m.actuator_actrange, m.actuator_biasprm, m.actuator_biastype, m.actuator_cranklength, m.actuator_ctrllimited, m.actuator_ctrlrange, m._impl.actuator_delay, m.actuator_dynprm, m.actuator_dyntype, m.actuator_forcelimited, m.actuator_forcerange, m.actuator_gainprm, m.actuator_gaintype, m.actuator_gear, m._impl.actuator_history, m._impl.actuator_historyadr, m.actuator_lengthrange, m.actuator_trnid, m.actuator_trntype, m._impl.actuator_trntype_body_adr, m._impl.block_dim, m._impl.body_branch_start, m._impl.body_branches, m.body_dofadr, m.body_dofnum, m._impl.body_fluid_ellipsoid, m.body_geomadr, m.body_geomnum, m.body_gravcomp, m.body_inertia, m.body_invweight0, m.body_ipos, m.body_iquat, m._impl.body_isdofancestor, m.body_jntadr, m.body_jntnum, m.body_mass, m.body_mocapid, m.body_parentid, m.body_pos, m.body_quat, m.body_rootid, m.body_subtreemass, m._impl.body_tree, m.body_treeid, m.body_weldid, m.cam_bodyid, m.cam_fovy, m.cam_intrinsic, m.cam_mat0, m.cam_mode, m.cam_pos, m.cam_pos0, m.cam_poscom0, m.cam_quat, m.cam_resolution, m.cam_sensorsize, m.cam_targetbodyid, m.dof_armature, m.dof_bodyid, m.dof_damping, m.dof_dampingpoly, m.dof_frictionloss, m.dof_invweight0, m.dof_jntid, m._impl.dof_length, m.dof_parentid, m.dof_solimp, m.dof_solref, m.dof_treeid, m._impl.dof_tri_col, m._impl.dof_tri_row, m._impl.eq_connect_adr, m.eq_data, m._impl.eq_flex_adr, m._impl.eq_jnt_adr, m.eq_obj1id, m.eq_obj2id, m.eq_objtype, m.eq_solimp, m.eq_solref, m._impl.eq_ten_adr, m.eq_type, m._impl.eq_wld_adr, m._impl.flex_bending, m._impl.flex_bendingadr, m._impl.flex_centered, m._impl.flex_conaffinity, m._impl.flex_condim, m._impl.flex_contype, m._impl.flex_damping, m._impl.flex_dim, m._impl.flex_edge, m._impl.flex_edgeadr, m._impl.flex_edgeflap, m._impl.flex_edgenum, m._impl.flex_elem, m._impl.flex_elemadr, m._impl.flex_elemdataadr, m._impl.flex_elemedge, m._impl.flex_elemedgeadr, m._impl.flex_elemnum, m._impl.flex_friction, m._impl.flex_gap, m._impl.flex_margin, m._impl.flex_priority, m._impl.flex_radius, m._impl.flex_shell, m._impl.flex_shelldataadr, m._impl.flex_shellnum, m._impl.flex_solimp, m._impl.flex_solmix, m._impl.flex_solref, m._impl.flex_stiffness, m._impl.flex_stiffnessadr, m._impl.flex_vert, m.flex_vertadr, m._impl.flex_vertbodyid, m._impl.flex_vertflexid, m.flex_vertnum, m._impl.flexedge_J_colind, m._impl.flexedge_J_rowadr, m._impl.flexedge_J_rownnz, m._impl.flexedge_invweight0, m._impl.flexedge_length0, m.geom_aabb, m.geom_bodyid, m.geom_conaffinity, m.geom_condim, m.geom_contype, jax.numpy.expand_dims(m.geom_dataid, 0), m.geom_fluid, m.geom_friction, m.geom_gap, m.geom_group, m.geom_margin, m.geom_matid, m._impl.geom_pair_type_count, m._impl.geom_plugin_index, m.geom_pos, m.geom_priority, m.geom_quat, m.geom_rbound, m.geom_rgba, m.geom_size, m.geom_solimp, m.geom_solmix, m.geom_solref, m.geom_type, m._impl.has_fluid, m._impl.has_sdf_geom, m.hfield_adr, m.hfield_data, m.hfield_ncol, m.hfield_nrow, m.hfield_size, m._impl.is_sparse, m.jnt_actfrclimited, m.jnt_actfrcrange, m.jnt_actgravcomp, m.jnt_axis, m.jnt_bodyid, m.jnt_dofadr, m._impl.jnt_limited_ball_adr, m._impl.jnt_limited_slide_hinge_adr, m.jnt_margin, m.jnt_pos, m.jnt_qposadr, m.jnt_range, m.jnt_solimp, m.jnt_solref, m.jnt_stiffness, m.jnt_stiffnesspoly, m.jnt_type, m._impl.light_bodyid, m.light_dir, m.light_dir0, m.light_mode, m.light_pos, m.light_pos0, m.light_poscom0, m._impl.light_targetbodyid, m._impl.mapM2D, m.mat_rgba, m._impl.max_ten_J_rownnz, m.mesh_face, m.mesh_faceadr, m.mesh_graph, m.mesh_graphadr, m.mesh_normal, m.mesh_normaladr, m.mesh_normalnum, m.mesh_octadr, m._impl.mesh_polyadr, m._impl.mesh_polymap, m._impl.mesh_polymapadr, m._impl.mesh_polymapnum, m._impl.mesh_polynormal, m._impl.mesh_polynum, m._impl.mesh_polyvert, m._impl.mesh_polyvertadr, m._impl.mesh_polyvertnum, m.mesh_quat, m.mesh_vert, m.mesh_vertadr, m.mesh_vertnum, m.nC, m.nD, m.nJten, m.na, m._impl.nacttrnbody, m.nbody, m._impl.nbranch, m.ncam, m.neq, m.nflex, m._impl.nflexedge, m._impl.nflexelem, m._impl.nflexshelldata, m._impl.nflexvert, m.ngeom, m.ngravcomp, m.nhistory, m.njnt, m.nlight, m._impl.nmaxcondim, m._impl.nmaxmeshdeg, m._impl.nmaxpolygon, m._impl.nmaxpyramid, m.nmeshface, m._impl.nrangefinder, m._impl.nsensorcollision, m._impl.nsensorcontact, m._impl.nsensortaxel, m.nsite, m.ntendon, m._impl.ntree, m.nu, m.nv, m._impl.nv_pad, m.nwrap, m._impl.nxn_geom_pair_filtered, m._impl.nxn_pairid, m._impl.nxn_pairid_filtered, m._impl.oct_aabb, m._impl.oct_child, m._impl.oct_coeff, m.pair_dim, m.pair_friction, m.pair_gap, m.pair_margin, m.pair_solimp, m.pair_solref, m.pair_solreffriction, m._impl.plugin, m._impl.plugin_attr, m._impl.qD_fullm_i, m._impl.qD_fullm_j, m._impl.qLD_all_updates, m._impl.qLD_level_offsets, m._impl.qLD_updates, m.qpos0, m.qpos_spring, m._impl.rangefinder_sensor_adr, m._impl.sensor_acc_adr, m.sensor_adr, m._impl.sensor_adr_to_contact_adr, m._impl.sensor_contact_adr, m.sensor_cutoff, m.sensor_datatype, m._impl.sensor_delay, m.sensor_dim, m._impl.sensor_e_kinetic, m._impl.sensor_e_potential, m._impl.sensor_history, m._impl.sensor_historyadr, m._impl.sensor_interval, m.sensor_intprm, m._impl.sensor_limitfrc_adr, m._impl.sensor_limitpos_adr, m._impl.sensor_limitvel_adr, m.sensor_objid, m.sensor_objtype, m._impl.sensor_pos_adr, m._impl.sensor_rangefinder_adr, m._impl.sensor_rangefinder_bodyid, m.sensor_refid, m.sensor_reftype, m._impl.sensor_rne_postconstraint, m._impl.sensor_subtree_vel, m._impl.sensor_tendonactfrc_adr, m._impl.sensor_touch_adr, m.sensor_type, m._impl.sensor_vel_adr, m.site_bodyid, m.site_pos, m.site_quat, m.site_size, m.site_type, m._impl.taxel_sensorid, m._impl.taxel_vertadr, m._impl.ten_J_colind, m._impl.ten_J_rowadr, m._impl.ten_J_rownnz, m.tendon_actfrclimited, m.tendon_actfrcrange, m.tendon_adr, m.tendon_armature, m.tendon_damping, m.tendon_dampingpoly, m.tendon_frictionloss, m._impl.tendon_geom_adr, m.tendon_invweight0, m._impl.tendon_jnt_adr, m.tendon_length0, m.tendon_lengthspring, m.tendon_limited, m._impl.tendon_limited_adr, m.tendon_margin, m.tendon_num, m.tendon_range, m._impl.tendon_site_pair_adr, m.tendon_solimp_fri, m.tendon_solimp_lim, m.tendon_solref_fri, m.tendon_solref_lim, m.tendon_stiffness, m.tendon_stiffnesspoly, m._impl.tree_dofadr, m._impl.tree_dofnum, m._impl.tree_sleep_policy, m._impl.wrap_geom_adr, m._impl.wrap_jnt_adr, m.wrap_objid, m.wrap_prm, m._impl.wrap_pulley_scale, m._impl.wrap_site_pair_adr, m.wrap_type, m.opt._impl.broadphase, m.opt._impl.broadphase_filter, m.opt._impl.ccd_iterations, m.opt._impl.ccd_tolerance, m.opt.cone, m.opt._impl.contact_sensor_maxmatch, m.opt.density, m.opt.disableflags, m.opt.enableflags, m.opt._impl.graph_conditional, m.opt.gravity, m.opt._impl.impratio_invsqrt, 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, m.opt._impl.sdf_initpoints, m.opt._impl.sdf_iterations, m.opt._impl.sleep_tolerance, m.opt.solver, m.opt.timestep, m.opt.tolerance, m.opt.viscosity, m.opt.wind, m.stat.meaninertia, d._impl.naccdmax, d._impl.naconmax, d._impl.njmax, d._impl.njmax_nnz, d._impl.M, d.act, d.act_dot, d.actuator_force, d.actuator_length, d._impl.actuator_moment, d._impl.actuator_velocity, d._impl.body_awake, d._impl.body_awake_ind, d._impl.cacc, d.cam_xmat, d.cam_xpos, d.cdof, d.cdof_dot, d._impl.cfrc_ext, d._impl.cfrc_int, d._impl.cinert, d._impl.crb, d.ctrl, d.cvel, d._impl.dof_awake_ind, d._impl.dof_island, d._impl.dof_islandid, d._impl.efc_islandid, d._impl.energy, d.eq_active, d._impl.flexedge_J, d._impl.flexedge_length, d._impl.flexedge_velocity, d._impl.flexvert_xpos, d.geom_xmat, d.geom_xpos, d.history, d._impl.iqacc, d._impl.iqacc_smooth, d._impl.iqfrc_constraint, d._impl.iqfrc_smooth, d._impl.island_dofadr, d._impl.island_efcadr, d._impl.island_idofadr, d._impl.island_ne, d._impl.island_nefc, d._impl.island_nf, d._impl.island_nv, d._impl.light_xdir, d._impl.light_xpos, d._impl.map_dof2idof, d._impl.map_efc2iefc, d._impl.map_idof2dof, d._impl.map_iefc2efc, d.mocap_pos, d.mocap_quat, d._impl.moment_colind, d._impl.moment_rowadr, d._impl.moment_rownnz, d._impl.nacon, d._impl.nbody_awake, d._impl.ncollision, d._impl.ne, d._impl.nefc, d._impl.nf, d._impl.nidof, d._impl.nisland, d._impl.nl, d._impl.ntree_awake, d._impl.nv_awake, d._impl.qLD, d._impl.qLDiagInv, d._impl.qLU, d.qacc, d.qacc_smooth, d.qacc_warmstart, d.qfrc_actuator, d.qfrc_applied, d.qfrc_bias, d.qfrc_constraint, d._impl.qfrc_damper, d.qfrc_fluid, d.qfrc_gravcomp, d.qfrc_passive, d.qfrc_smooth, d._impl.qfrc_spring, d.qpos, d.qvel, d.sensordata, d.site_xmat, d.site_xpos, d._impl.solver_niter, d._impl.subtree_angmom, d.subtree_com, d._impl.subtree_linvel, d._impl.ten_J, d.ten_length, d._impl.ten_velocity, d._impl.ten_wrapadr, d._impl.ten_wrapnum, d.time, d._impl.tree_asleep, d._impl.tree_awake, d._impl.tree_island, d._impl.wrap_obj, d._impl.wrap_xpos, d.xanchor, d.xaxis, d.xfrc_applied, d.ximat, d.xipos, d.xmat, d.xpos, d.xquat, d._impl.contact__dim, d._impl.contact__dist, d._impl.contact__efc_address, d._impl.contact__flex, d._impl.contact__frame, d._impl.contact__friction, d._impl.contact__geom, d._impl.contact__geomcollisionid, d._impl.contact__includemargin, d._impl.contact__pos, d._impl.contact__solimp, d._impl.contact__solref, d._impl.contact__solreffriction, d._impl.contact__type, d._impl.contact__vert, d._impl.contact__worldid, d._impl.efc__D, d._impl.efc__J, d._impl.efc__J_colind, d._impl.efc__J_rowadr, d._impl.efc__J_rownnz, d._impl.efc__Jqvel, d._impl.efc__Ma, d._impl.efc__aref, d._impl.efc__force, d._impl.efc__frictionloss, d._impl.efc__iD, d._impl.efc__iJ, d._impl.efc__iJ_colind, d._impl.efc__iJ_rowadr, d._impl.efc__iJ_rownnz, d._impl.efc__iaref, d._impl.efc__id, d._impl.efc__iforce, d._impl.efc__ifrictionloss, d._impl.efc__iid, d._impl.efc__island, d._impl.efc__istate, d._impl.efc__itype, d._impl.efc__margin, d._impl.efc__pos, d._impl.efc__state, d._impl.efc__type, d._impl.efc__vel, ) d = d.tree_replace({ '_impl.M': out[0], 'act': out[1], 'act_dot': out[2], 'actuator_force': out[3], 'actuator_length': out[4], '_impl.actuator_moment': out[5], '_impl.actuator_velocity': out[6], '_impl.body_awake': out[7], '_impl.body_awake_ind': out[8], '_impl.cacc': out[9], 'cam_xmat': out[10], 'cam_xpos': out[11], 'cdof': out[12], 'cdof_dot': out[13], '_impl.cfrc_ext': out[14], '_impl.cfrc_int': out[15], '_impl.cinert': out[16], '_impl.crb': out[17], 'cvel': out[18], '_impl.dof_awake_ind': out[19], '_impl.dof_island': out[20], '_impl.dof_islandid': out[21], '_impl.efc_islandid': out[22], '_impl.energy': out[23], '_impl.flexedge_J': out[24], '_impl.flexedge_length': out[25], '_impl.flexedge_velocity': out[26], '_impl.flexvert_xpos': out[27], 'geom_xmat': out[28], 'geom_xpos': out[29], 'history': out[30], '_impl.iqacc': out[31], '_impl.iqacc_smooth': out[32], '_impl.iqfrc_constraint': out[33], '_impl.iqfrc_smooth': out[34], '_impl.island_dofadr': out[35], '_impl.island_efcadr': out[36], '_impl.island_idofadr': out[37], '_impl.island_ne': out[38], '_impl.island_nefc': out[39], '_impl.island_nf': out[40], '_impl.island_nv': out[41], '_impl.light_xdir': out[42], '_impl.light_xpos': out[43], '_impl.map_dof2idof': out[44], '_impl.map_efc2iefc': out[45], '_impl.map_idof2dof': out[46], '_impl.map_iefc2efc': out[47], '_impl.moment_colind': out[48], '_impl.moment_rowadr': out[49], '_impl.moment_rownnz': out[50], '_impl.nacon': out[51], '_impl.nbody_awake': out[52], '_impl.ncollision': out[53], '_impl.ne': out[54], '_impl.nefc': out[55], '_impl.nf': out[56], '_impl.nidof': out[57], '_impl.nisland': out[58], '_impl.nl': out[59], '_impl.ntree_awake': out[60], '_impl.nv_awake': out[61], '_impl.qLD': out[62], '_impl.qLDiagInv': out[63], '_impl.qLU': out[64], 'qacc': out[65], 'qacc_smooth': out[66], 'qacc_warmstart': out[67], 'qfrc_actuator': out[68], 'qfrc_bias': out[69], 'qfrc_constraint': out[70], '_impl.qfrc_damper': out[71], 'qfrc_fluid': out[72], 'qfrc_gravcomp': out[73], 'qfrc_passive': out[74], 'qfrc_smooth': out[75], '_impl.qfrc_spring': out[76], 'qpos': out[77], 'qvel': out[78], 'sensordata': out[79], 'site_xmat': out[80], 'site_xpos': out[81], '_impl.solver_niter': out[82], '_impl.subtree_angmom': out[83], 'subtree_com': out[84], '_impl.subtree_linvel': out[85], '_impl.ten_J': out[86], 'ten_length': out[87], '_impl.ten_velocity': out[88], '_impl.ten_wrapadr': out[89], '_impl.ten_wrapnum': out[90], 'time': out[91], '_impl.tree_asleep': out[92], '_impl.tree_awake': out[93], '_impl.tree_island': out[94], '_impl.wrap_obj': out[95], '_impl.wrap_xpos': out[96], 'xanchor': out[97], 'xaxis': out[98], 'ximat': out[99], 'xipos': out[100], 'xmat': out[101], 'xpos': out[102], 'xquat': out[103], '_impl.contact__dim': out[104], '_impl.contact__dist': out[105], '_impl.contact__efc_address': out[106], '_impl.contact__flex': out[107], '_impl.contact__frame': out[108], '_impl.contact__friction': out[109], '_impl.contact__geom': out[110], '_impl.contact__geomcollisionid': out[111], '_impl.contact__includemargin': out[112], '_impl.contact__pos': out[113], '_impl.contact__solimp': out[114], '_impl.contact__solref': out[115], '_impl.contact__solreffriction': out[116], '_impl.contact__type': out[117], '_impl.contact__vert': out[118], '_impl.contact__worldid': out[119], '_impl.efc__D': out[120], '_impl.efc__J': out[121], '_impl.efc__J_colind': out[122], '_impl.efc__J_rowadr': out[123], '_impl.efc__J_rownnz': out[124], '_impl.efc__Jqvel': out[125], '_impl.efc__Ma': out[126], '_impl.efc__aref': out[127], '_impl.efc__force': out[128], '_impl.efc__frictionloss': out[129], '_impl.efc__iD': out[130], '_impl.efc__iJ': out[131], '_impl.efc__iJ_colind': out[132], '_impl.efc__iJ_rowadr': out[133], '_impl.efc__iJ_rownnz': out[134], '_impl.efc__iaref': out[135], '_impl.efc__id': out[136], '_impl.efc__iforce': out[137], '_impl.efc__ifrictionloss': out[138], '_impl.efc__iid': out[139], '_impl.efc__island': out[140], '_impl.efc__istate': out[141], '_impl.efc__itype': out[142], '_impl.efc__margin': out[143], '_impl.efc__pos': out[144], '_impl.efc__state': out[145], '_impl.efc__type': out[146], '_impl.efc__vel': out[147], }) return d @jax.custom_batching.custom_vmap @ffi.marshal_jax_warp_callable def step(m: types.Model, d: types.Data): return _step_jax_impl(m, d) @step.def_vmap @ffi.marshal_custom_vmap def step_vmap(unused_axis_size, is_batched, m: types.Model, d: types.Data): d = step(m, d) return d, is_batched[1]