# Copyright 2025 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.
# ==============================================================================
"""MJX Warp types.
DO NOT EDIT. This file is auto-generated.
"""
import dataclasses
import typing
from typing import Tuple
import jax
from jax import tree_util
from jax.interpreters import batching
from mujoco.mjx._src import dataclasses as mjx_dataclasses
import numpy as np
if typing.TYPE_CHECKING:
  GraphMode = int
else:
  try:
    from warp._src.jax_experimental.ffi import GraphMode
  except ImportError:
    GraphMode = int
PyTreeNode = mjx_dataclasses.PyTreeNode

@dataclasses.dataclass(frozen=True)
@tree_util.register_pytree_node_class
class TileSet:
  """Tiling configuration for decomposable block diagonal matrix.

  For non-square, non-block-diagonal tiles, use two tilesets.

  Attributes:
    adr: address of each tile in the set
    size: size of all the tiles in this set
  """

  adr: np.ndarray
  size: int

  def tree_flatten(self):
    children = list((getattr(self, k) for k in self.__dataclass_fields__))
    return (children, None)

  @classmethod
  def tree_unflatten(cls, aux_data, children):
    del aux_data
    return cls(*children)


@dataclasses.dataclass(frozen=True)
@tree_util.register_pytree_node_class
class BlockDim:
  """Block dimension 'block_dim' settings for wp.launch_tiled.

  TODO(team): experimental and may be removed
  """
  actuator_velocity: int
  cholesky_factorize: int
  cholesky_factorize_solve: int
  cholesky_solve: int
  contact_sort: int
  energy_vel_kinetic: int
  euler_dense: int
  linesearch_iterative: int
  qderiv_actuator_dense: int
  ray: int
  segmented_sort: int
  tendon_velocity: int
  update_gradient_JTDAJ_dense: int
  update_gradient_JTDAJ_sparse: int
  update_gradient_cholesky: int
  update_gradient_cholesky_blocked: int

  def tree_flatten(self):
    children = list((getattr(self, k) for k in self.__dataclass_fields__))
    return (children, None)

  @classmethod
  def tree_unflatten(cls, aux_data, children):
    del aux_data
    return cls(*children)


class StatisticWarp(PyTreeNode):
  """Derived fields from Statistic."""
  meaninertia: jax.Array

class OptionWarp(PyTreeNode):
  """Derived fields from Option."""
  broadphase: int
  broadphase_filter: int
  ccd_iterations: int
  ccd_tolerance: jax.Array
  contact_sensor_maxmatch: int
  graph_conditional: bool
  graph_mode: GraphMode
  impratio_invsqrt: jax.Array
  ls_parallel: bool
  ls_parallel_min_step: float
  run_collision_detection: bool
  sdf_initpoints: int
  sdf_iterations: int

class ModelWarp(PyTreeNode):
  """Derived fields from Model."""
  M_colind: np.ndarray
  M_rowadr: np.ndarray
  M_rownnz: np.ndarray
  actuator_trntype_body_adr: np.ndarray
  block_dim: BlockDim
  body_branch_start: np.ndarray
  body_branches: np.ndarray
  body_fluid_ellipsoid: np.ndarray
  body_tree: Tuple[np.ndarray, ...]
  cam_projection: np.ndarray
  collision_sensor_adr: np.ndarray
  dof_tri_col: np.ndarray
  dof_tri_row: np.ndarray
  eq_connect_adr: np.ndarray
  eq_flex_adr: np.ndarray
  eq_jnt_adr: np.ndarray
  eq_ten_adr: np.ndarray
  eq_wld_adr: np.ndarray
  flex_bending: np.ndarray
  flex_damping: np.ndarray
  flex_dim: np.ndarray
  flex_edge: np.ndarray
  flex_edgeadr: np.ndarray
  flex_edgeflap: np.ndarray
  flex_edgenum: np.ndarray
  flex_elem: np.ndarray
  flex_elemadr: np.ndarray
  flex_elemedge: np.ndarray
  flex_elemedgeadr: np.ndarray
  flex_elemnum: np.ndarray
  flex_stiffness: np.ndarray
  flex_vertadr: np.ndarray
  flex_vertbodyid: np.ndarray
  flex_vertnum: np.ndarray
  flexedge_J_colind: np.ndarray
  flexedge_J_rowadr: np.ndarray
  flexedge_J_rownnz: np.ndarray
  flexedge_invweight0: np.ndarray
  flexedge_length0: np.ndarray
  geom_pair_type_count: Tuple[int, ...]
  geom_plugin_index: np.ndarray
  has_fluid: bool
  has_sdf_geom: bool
  is_sparse: bool
  jnt_limited_ball_adr: np.ndarray
  jnt_limited_slide_hinge_adr: np.ndarray
  light_active: jax.Array
  light_bodyid: np.ndarray
  light_targetbodyid: np.ndarray
  mapM2M: np.ndarray
  mat_texrepeat: jax.Array
  mesh_polyadr: np.ndarray
  mesh_polymap: np.ndarray
  mesh_polymapadr: np.ndarray
  mesh_polymapnum: np.ndarray
  mesh_polynormal: np.ndarray
  mesh_polynum: np.ndarray
  mesh_polyvert: np.ndarray
  mesh_polyvertadr: np.ndarray
  mesh_polyvertnum: np.ndarray
  mocap_bodyid: np.ndarray
  nacttrnbody: int
  nbranch: int
  nflex: int
  nflexedge: int
  nflexelem: int
  nflexelemdata: int
  nflexelemedge: int
  nflexvert: int
  nmaxcondim: int
  nmaxmeshdeg: int
  nmaxpolygon: int
  nmaxpyramid: int
  noct: int
  nplugin: int
  nrangefinder: int
  nsensorcollision: int
  nsensorcontact: int
  nsensortaxel: int
  ntree: int
  nv_pad: int
  nxn_geom_pair: np.ndarray
  nxn_geom_pair_filtered: np.ndarray
  nxn_pairid: np.ndarray
  nxn_pairid_filtered: np.ndarray
  oct_aabb: np.ndarray
  oct_child: np.ndarray
  oct_coeff: np.ndarray
  plugin: np.ndarray
  plugin_attr: np.ndarray
  qLD_updates: Tuple[np.ndarray, ...]
  qM_fullm_i: np.ndarray
  qM_fullm_j: np.ndarray
  qM_mulm_col: np.ndarray
  qM_mulm_madr: np.ndarray
  qM_mulm_rowadr: np.ndarray
  qM_tiles: Tuple[TileSet, ...]
  rangefinder_sensor_adr: np.ndarray
  sensor_acc_adr: np.ndarray
  sensor_adr_to_contact_adr: np.ndarray
  sensor_collision_start_adr: np.ndarray
  sensor_contact_adr: np.ndarray
  sensor_e_kinetic: bool
  sensor_e_potential: bool
  sensor_limitfrc_adr: np.ndarray
  sensor_limitpos_adr: np.ndarray
  sensor_limitvel_adr: np.ndarray
  sensor_pos_adr: np.ndarray
  sensor_rangefinder_adr: np.ndarray
  sensor_rangefinder_bodyid: np.ndarray
  sensor_rne_postconstraint: bool
  sensor_subtree_vel: bool
  sensor_tendonactfrc_adr: np.ndarray
  sensor_touch_adr: np.ndarray
  sensor_vel_adr: np.ndarray
  taxel_sensorid: np.ndarray
  taxel_vertadr: np.ndarray
  ten_wrapadr_site: np.ndarray
  ten_wrapnum_site: np.ndarray
  tendon_geom_adr: np.ndarray
  tendon_jnt_adr: np.ndarray
  tendon_limited_adr: np.ndarray
  tendon_site_pair_adr: np.ndarray
  tree_bodynum: np.ndarray
  tree_dofadr: np.ndarray
  tree_dofnum: np.ndarray
  wrap_geom_adr: np.ndarray
  wrap_jnt_adr: np.ndarray
  wrap_pulley_scale: np.ndarray
  wrap_site_adr: np.ndarray
  wrap_site_pair_adr: np.ndarray

class DataWarp(PyTreeNode):
  """Derived fields from Data."""
  actuator_moment: jax.Array
  actuator_velocity: jax.Array
  cacc: jax.Array
  cfrc_ext: jax.Array
  cfrc_int: jax.Array
  cinert: jax.Array
  contact__dim: jax.Array
  contact__dist: jax.Array
  contact__efc_address: jax.Array
  contact__frame: jax.Array
  contact__friction: jax.Array
  contact__geom: jax.Array
  contact__geomcollisionid: jax.Array
  contact__includemargin: jax.Array
  contact__pos: jax.Array
  contact__solimp: jax.Array
  contact__solref: jax.Array
  contact__solreffriction: jax.Array
  contact__type: jax.Array
  contact__worldid: jax.Array
  crb: jax.Array
  efc__D: jax.Array
  efc__J: jax.Array
  efc__Ma: jax.Array
  efc__aref: jax.Array
  efc__force: jax.Array
  efc__frictionloss: jax.Array
  efc__id: jax.Array
  efc__margin: jax.Array
  efc__pos: jax.Array
  efc__state: jax.Array
  efc__type: jax.Array
  efc__vel: jax.Array
  energy: jax.Array
  flexedge_J: jax.Array
  flexedge_length: jax.Array
  flexedge_velocity: jax.Array
  flexvert_xpos: jax.Array
  light_xdir: jax.Array
  light_xpos: jax.Array
  naccdmax: int
  nacon: jax.Array
  naconmax: int
  ncollision: jax.Array
  ne: jax.Array
  nefc: jax.Array
  nf: jax.Array
  njmax: int
  nl: jax.Array
  nworld: int
  qLD: jax.Array
  qLDiagInv: jax.Array
  qM: jax.Array
  qfrc_damper: jax.Array
  qfrc_spring: jax.Array
  solver_niter: jax.Array
  subtree_angmom: jax.Array
  subtree_linvel: jax.Array
  ten_J: jax.Array
  ten_velocity: jax.Array
  ten_wrapadr: jax.Array
  ten_wrapnum: jax.Array
  wrap_obj: jax.Array
  wrap_xpos: jax.Array
  shape = property(lambda self: self.cacc.shape)
DATA_NON_VMAP = {
    'contact__dim',
    'contact__dist',
    'contact__efc_address',
    'contact__frame',
    'contact__friction',
    'contact__geom',
    'contact__geomcollisionid',
    'contact__includemargin',
    'contact__pos',
    'contact__solimp',
    'contact__solref',
    'contact__solreffriction',
    'contact__type',
    'contact__worldid',
    'naccdmax',
    'nacon',
    'naconmax',
    'ncollision',
    'njmax',
    'nworld',
}

def _to_elt(cont, _, d, axis):
  return DataWarp(**{
      f.name: (
          cont(getattr(d, f.name), axis)
          if f.name not in DATA_NON_VMAP
          else getattr(d, f.name)
      )
      for f in DataWarp.fields()
  })


def _from_elt(cont, axis_size, d, axis_dest):
  return DataWarp(**{
      f.name: (
          cont(axis_size, getattr(d, f.name), axis_dest)
          if f.name not in DATA_NON_VMAP
          else getattr(d, f.name)
      )
      for f in DataWarp.fields()
  })


batching.register_vmappable(DataWarp, int, int, _to_elt, _from_elt, None)

_NDIM = {
    'Data': {
        'act': 2,
        'act_dot': 2,
        'actuator_force': 2,
        'actuator_length': 2,
        'actuator_moment': 3,
        'actuator_velocity': 2,
        'cacc': 3,
        'cam_xmat': 4,
        'cam_xpos': 3,
        'cdof': 3,
        'cdof_dot': 3,
        'cfrc_ext': 3,
        'cfrc_int': 3,
        'cinert': 3,
        'contact__dim': 1,
        'contact__dist': 1,
        'contact__efc_address': 2,
        'contact__frame': 3,
        'contact__friction': 2,
        'contact__geom': 2,
        'contact__geomcollisionid': 1,
        'contact__includemargin': 1,
        'contact__pos': 2,
        'contact__solimp': 2,
        'contact__solref': 2,
        'contact__solreffriction': 2,
        'contact__type': 1,
        'contact__worldid': 1,
        'crb': 3,
        'ctrl': 2,
        'cvel': 3,
        'efc__D': 2,
        'efc__J': 3,
        'efc__Ma': 2,
        'efc__aref': 2,
        'efc__force': 2,
        'efc__frictionloss': 2,
        'efc__id': 2,
        'efc__margin': 2,
        'efc__pos': 2,
        'efc__state': 2,
        'efc__type': 2,
        'efc__vel': 2,
        'energy': 2,
        'eq_active': 2,
        'flexedge_J': 3,
        'flexedge_length': 2,
        'flexedge_velocity': 2,
        'flexvert_xpos': 3,
        'geom_xmat': 4,
        'geom_xpos': 3,
        'light_xdir': 3,
        'light_xpos': 3,
        'mocap_pos': 3,
        'mocap_quat': 3,
        'naccdmax': 0,
        'nacon': 1,
        'naconmax': 0,
        'ncollision': 1,
        'ne': 1,
        'nefc': 1,
        'nf': 1,
        'njmax': 0,
        'nl': 1,
        'nworld': 0,
        'qLD': 3,
        'qLDiagInv': 2,
        'qM': 3,
        'qacc': 2,
        'qacc_smooth': 2,
        'qacc_warmstart': 2,
        'qfrc_actuator': 2,
        'qfrc_applied': 2,
        'qfrc_bias': 2,
        'qfrc_constraint': 2,
        'qfrc_damper': 2,
        'qfrc_fluid': 2,
        'qfrc_gravcomp': 2,
        'qfrc_inverse': 2,
        'qfrc_passive': 2,
        'qfrc_smooth': 2,
        'qfrc_spring': 2,
        'qpos': 2,
        'qvel': 2,
        'sensordata': 2,
        'site_xmat': 4,
        'site_xpos': 3,
        'solver_niter': 1,
        'subtree_angmom': 3,
        'subtree_com': 3,
        'subtree_linvel': 3,
        'ten_J': 3,
        'ten_length': 2,
        'ten_velocity': 2,
        'ten_wrapadr': 2,
        'ten_wrapnum': 2,
        'time': 1,
        'wrap_obj': 3,
        'wrap_xpos': 3,
        'xanchor': 3,
        'xaxis': 3,
        'xfrc_applied': 3,
        'ximat': 4,
        'xipos': 3,
        'xmat': 4,
        'xpos': 3,
        'xquat': 3,
    },
    'Model': {
        'M_colind': 1,
        'M_rowadr': 1,
        'M_rownnz': 1,
        'actuator_acc0': 1,
        'actuator_actadr': 1,
        'actuator_actearly': 1,
        'actuator_actlimited': 1,
        'actuator_actnum': 1,
        'actuator_actrange': 3,
        'actuator_biasprm': 3,
        'actuator_biastype': 1,
        'actuator_cranklength': 1,
        'actuator_ctrllimited': 1,
        'actuator_ctrlrange': 3,
        'actuator_dynprm': 3,
        'actuator_dyntype': 1,
        'actuator_forcelimited': 1,
        'actuator_forcerange': 3,
        'actuator_gainprm': 3,
        'actuator_gaintype': 1,
        'actuator_gear': 3,
        'actuator_lengthrange': 2,
        'actuator_trnid': 2,
        'actuator_trntype': 1,
        'actuator_trntype_body_adr': 1,
        'block_dim__actuator_velocity': 0,
        'block_dim__cholesky_factorize': 0,
        'block_dim__cholesky_factorize_solve': 0,
        'block_dim__cholesky_solve': 0,
        'block_dim__contact_sort': 0,
        'block_dim__energy_vel_kinetic': 0,
        'block_dim__euler_dense': 0,
        'block_dim__linesearch_iterative': 0,
        'block_dim__qderiv_actuator_dense': 0,
        'block_dim__ray': 0,
        'block_dim__segmented_sort': 0,
        'block_dim__tendon_velocity': 0,
        'block_dim__update_gradient_JTDAJ_dense': 0,
        'block_dim__update_gradient_JTDAJ_sparse': 0,
        'block_dim__update_gradient_cholesky': 0,
        'block_dim__update_gradient_cholesky_blocked': 0,
        'body_branch_start': 1,
        'body_branches': 1,
        'body_conaffinity': 1,
        'body_contype': 1,
        'body_dofadr': 1,
        'body_dofnum': 1,
        'body_fluid_ellipsoid': 1,
        'body_geomadr': 1,
        'body_geomnum': 1,
        'body_gravcomp': 2,
        'body_inertia': 3,
        'body_invweight0': 3,
        'body_ipos': 3,
        'body_iquat': 3,
        'body_jntadr': 1,
        'body_jntnum': 1,
        'body_mass': 2,
        'body_mocapid': 1,
        'body_parentid': 1,
        'body_pos': 3,
        'body_quat': 3,
        'body_rootid': 1,
        'body_subtreemass': 2,
        'body_tree': -1,
        'body_treeid': 1,
        'body_weldid': 1,
        'cam_bodyid': 1,
        'cam_fovy': 2,
        'cam_intrinsic': 3,
        'cam_mat0': 4,
        'cam_mode': 1,
        'cam_pos': 3,
        'cam_pos0': 3,
        'cam_poscom0': 3,
        'cam_projection': 1,
        'cam_quat': 3,
        'cam_resolution': 2,
        'cam_sensorsize': 2,
        'cam_targetbodyid': 1,
        'collision_sensor_adr': 1,
        'dof_Madr': 1,
        'dof_armature': 2,
        'dof_bodyid': 1,
        'dof_damping': 2,
        'dof_frictionloss': 2,
        'dof_invweight0': 2,
        'dof_jntid': 1,
        'dof_parentid': 1,
        'dof_solimp': 3,
        'dof_solref': 3,
        'dof_treeid': 1,
        'dof_tri_col': 1,
        'dof_tri_row': 1,
        'eq_active0': 1,
        'eq_connect_adr': 1,
        'eq_data': 3,
        'eq_flex_adr': 1,
        'eq_jnt_adr': 1,
        'eq_obj1id': 1,
        'eq_obj2id': 1,
        'eq_objtype': 1,
        'eq_solimp': 3,
        'eq_solref': 3,
        'eq_ten_adr': 1,
        'eq_type': 1,
        'eq_wld_adr': 1,
        'exclude_signature': 1,
        'flex_bending': 2,
        'flex_damping': 1,
        'flex_dim': 1,
        'flex_edge': 2,
        'flex_edgeadr': 1,
        'flex_edgeflap': 2,
        'flex_edgenum': 1,
        'flex_elem': 1,
        'flex_elemadr': 1,
        'flex_elemedge': 1,
        'flex_elemedgeadr': 1,
        'flex_elemnum': 1,
        'flex_stiffness': 2,
        'flex_vertadr': 1,
        'flex_vertbodyid': 1,
        'flex_vertnum': 1,
        'flexedge_J_colind': 1,
        'flexedge_J_rowadr': 1,
        'flexedge_J_rownnz': 1,
        'flexedge_invweight0': 1,
        'flexedge_length0': 1,
        'geom_aabb': 4,
        'geom_bodyid': 1,
        'geom_conaffinity': 1,
        'geom_condim': 1,
        'geom_contype': 1,
        'geom_dataid': 1,
        'geom_fluid': 2,
        'geom_friction': 3,
        'geom_gap': 2,
        'geom_group': 1,
        'geom_margin': 2,
        'geom_matid': 2,
        'geom_pair_type_count': -1,
        'geom_plugin_index': 1,
        'geom_pos': 3,
        'geom_priority': 1,
        'geom_quat': 3,
        'geom_rbound': 2,
        'geom_rgba': 3,
        'geom_size': 3,
        'geom_solimp': 3,
        'geom_solmix': 2,
        'geom_solref': 3,
        'geom_type': 1,
        'has_fluid': 0,
        'has_sdf_geom': 0,
        'hfield_adr': 1,
        'hfield_data': 1,
        'hfield_ncol': 1,
        'hfield_nrow': 1,
        'hfield_size': 2,
        'is_sparse': 0,
        'jnt_actfrclimited': 1,
        'jnt_actfrcrange': 3,
        'jnt_actgravcomp': 1,
        'jnt_axis': 3,
        'jnt_bodyid': 1,
        'jnt_dofadr': 1,
        'jnt_limited': 1,
        'jnt_limited_ball_adr': 1,
        'jnt_limited_slide_hinge_adr': 1,
        'jnt_margin': 2,
        'jnt_pos': 3,
        'jnt_qposadr': 1,
        'jnt_range': 3,
        'jnt_solimp': 3,
        'jnt_solref': 3,
        'jnt_stiffness': 2,
        'jnt_type': 1,
        'light_active': 2,
        'light_bodyid': 1,
        'light_castshadow': 2,
        'light_dir': 3,
        'light_dir0': 3,
        'light_mode': 1,
        'light_pos': 3,
        'light_pos0': 3,
        'light_poscom0': 3,
        'light_targetbodyid': 1,
        'light_type': 2,
        'mapM2M': 1,
        'mat_rgba': 3,
        'mat_texid': 3,
        'mat_texrepeat': 3,
        'mesh_face': 2,
        'mesh_faceadr': 1,
        'mesh_graph': 1,
        'mesh_graphadr': 1,
        'mesh_normal': 2,
        'mesh_normaladr': 1,
        'mesh_polyadr': 1,
        'mesh_polymap': 1,
        'mesh_polymapadr': 1,
        'mesh_polymapnum': 1,
        'mesh_polynormal': 2,
        'mesh_polynum': 1,
        'mesh_polyvert': 1,
        'mesh_polyvertadr': 1,
        'mesh_polyvertnum': 1,
        'mesh_quat': 2,
        'mesh_vert': 2,
        'mesh_vertadr': 1,
        'mesh_vertnum': 1,
        'mocap_bodyid': 1,
        'nC': 0,
        'nM': 0,
        'na': 0,
        'nacttrnbody': 0,
        'nbody': 0,
        'nbranch': 0,
        'ncam': 0,
        'neq': 0,
        'nexclude': 0,
        'nflex': 0,
        'nflexedge': 0,
        'nflexelem': 0,
        'nflexelemdata': 0,
        'nflexelemedge': 0,
        'nflexvert': 0,
        'ngeom': 0,
        'ngravcomp': 0,
        'nhfield': 0,
        'nhfielddata': 0,
        'njnt': 0,
        'nlight': 0,
        'nmat': 0,
        'nmaxcondim': 0,
        'nmaxmeshdeg': 0,
        'nmaxpolygon': 0,
        'nmaxpyramid': 0,
        'nmesh': 0,
        'nmeshface': 0,
        'nmeshgraph': 0,
        'nmeshnormal': 0,
        'nmeshpoly': 0,
        'nmeshpolymap': 0,
        'nmeshpolyvert': 0,
        'nmeshvert': 0,
        'nmocap': 0,
        'noct': 0,
        'npair': 0,
        'nplugin': 0,
        'nq': 0,
        'nrangefinder': 0,
        'nsensor': 0,
        'nsensorcollision': 0,
        'nsensorcontact': 0,
        'nsensordata': 0,
        'nsensortaxel': 0,
        'nsite': 0,
        'ntendon': 0,
        'ntree': 0,
        'nu': 0,
        'nv': 0,
        'nv_pad': 0,
        'nwrap': 0,
        'nxn_geom_pair': 2,
        'nxn_geom_pair_filtered': 2,
        'nxn_pairid': 2,
        'nxn_pairid_filtered': 2,
        'oct_aabb': 3,
        'oct_child': 2,
        'oct_coeff': 2,
        'opt__ccd_iterations': 0,
        'opt__ccd_tolerance': 1,
        'opt__cone': 0,
        'opt__contact_sensor_maxmatch': 0,
        'opt__density': 1,
        'opt__disableflags': 0,
        'opt__enableflags': 0,
        'opt__graph_conditional': 0,
        'opt__gravity': 2,
        'opt__impratio_invsqrt': 1,
        'opt__integrator': 0,
        'opt__iterations': 0,
        'opt__ls_iterations': 0,
        'opt__ls_parallel': 0,
        'opt__ls_parallel_min_step': 0,
        'opt__ls_tolerance': 1,
        'opt__magnetic': 2,
        'opt__run_collision_detection': 0,
        'opt__sdf_initpoints': 0,
        'opt__sdf_iterations': 0,
        'opt__solver': 0,
        'opt__timestep': 1,
        'opt__tolerance': 1,
        'opt__viscosity': 1,
        'opt__wind': 2,
        'pair_dim': 1,
        'pair_friction': 3,
        'pair_gap': 2,
        'pair_geom1': 1,
        'pair_geom2': 1,
        'pair_margin': 2,
        'pair_solimp': 3,
        'pair_solref': 3,
        'pair_solreffriction': 3,
        'plugin': 1,
        'plugin_attr': 2,
        'qLD_updates': -1,
        'qM_fullm_i': 1,
        'qM_fullm_j': 1,
        'qM_mulm_col': 1,
        'qM_mulm_madr': 1,
        'qM_mulm_rowadr': 1,
        'qM_tiles': -1,
        'qpos0': 2,
        'qpos_spring': 2,
        'rangefinder_sensor_adr': 1,
        'sensor_acc_adr': 1,
        'sensor_adr': 1,
        'sensor_adr_to_contact_adr': 1,
        'sensor_collision_start_adr': 1,
        'sensor_contact_adr': 1,
        'sensor_cutoff': 1,
        'sensor_datatype': 1,
        'sensor_dim': 1,
        'sensor_e_kinetic': 0,
        'sensor_e_potential': 0,
        'sensor_intprm': 2,
        'sensor_limitfrc_adr': 1,
        'sensor_limitpos_adr': 1,
        'sensor_limitvel_adr': 1,
        'sensor_objid': 1,
        'sensor_objtype': 1,
        'sensor_pos_adr': 1,
        'sensor_rangefinder_adr': 1,
        'sensor_rangefinder_bodyid': 1,
        'sensor_refid': 1,
        'sensor_reftype': 1,
        'sensor_rne_postconstraint': 0,
        'sensor_subtree_vel': 0,
        'sensor_tendonactfrc_adr': 1,
        'sensor_touch_adr': 1,
        'sensor_type': 1,
        'sensor_vel_adr': 1,
        'site_bodyid': 1,
        'site_pos': 3,
        'site_quat': 3,
        'site_size': 2,
        'site_type': 1,
        'stat__meaninertia': 1,
        'taxel_sensorid': 1,
        'taxel_vertadr': 1,
        'ten_wrapadr_site': 1,
        'ten_wrapnum_site': 1,
        'tendon_actfrclimited': 1,
        'tendon_actfrcrange': 3,
        'tendon_adr': 1,
        'tendon_armature': 2,
        'tendon_damping': 2,
        'tendon_frictionloss': 2,
        'tendon_geom_adr': 1,
        'tendon_invweight0': 2,
        'tendon_jnt_adr': 1,
        'tendon_length0': 2,
        'tendon_lengthspring': 3,
        'tendon_limited': 1,
        'tendon_limited_adr': 1,
        'tendon_margin': 2,
        'tendon_num': 1,
        'tendon_range': 3,
        'tendon_site_pair_adr': 1,
        'tendon_solimp_fri': 3,
        'tendon_solimp_lim': 3,
        'tendon_solref_fri': 3,
        'tendon_solref_lim': 3,
        'tendon_stiffness': 2,
        'tree_bodynum': 1,
        'tree_dofadr': 1,
        'tree_dofnum': 1,
        'wrap_geom_adr': 1,
        'wrap_jnt_adr': 1,
        'wrap_objid': 1,
        'wrap_prm': 1,
        'wrap_pulley_scale': 1,
        'wrap_site_adr': 1,
        'wrap_site_pair_adr': 1,
        'wrap_type': 1,
    },
    'Option': {
        'ccd_iterations': 0,
        'ccd_tolerance': 1,
        'cone': 0,
        'contact_sensor_maxmatch': 0,
        'density': 1,
        'disableflags': 0,
        'enableflags': 0,
        'graph_conditional': 0,
        'gravity': 2,
        'impratio_invsqrt': 1,
        'integrator': 0,
        'iterations': 0,
        'ls_iterations': 0,
        'ls_parallel': 0,
        'ls_parallel_min_step': 0,
        'ls_tolerance': 1,
        'magnetic': 2,
        'run_collision_detection': 0,
        'sdf_initpoints': 0,
        'sdf_iterations': 0,
        'solver': 0,
        'timestep': 1,
        'tolerance': 1,
        'viscosity': 1,
        'wind': 2,
    },
    'Statistic': {'meaninertia': 1},
}
_BATCH_DIM = {
    'Data': {
        'act': True,
        'act_dot': True,
        'actuator_force': True,
        'actuator_length': True,
        'actuator_moment': True,
        'actuator_velocity': True,
        'cacc': True,
        'cam_xmat': True,
        'cam_xpos': True,
        'cdof': True,
        'cdof_dot': True,
        'cfrc_ext': True,
        'cfrc_int': True,
        'cinert': True,
        'contact__dim': False,
        'contact__dist': False,
        'contact__efc_address': False,
        'contact__frame': False,
        'contact__friction': False,
        'contact__geom': False,
        'contact__geomcollisionid': False,
        'contact__includemargin': False,
        'contact__pos': False,
        'contact__solimp': False,
        'contact__solref': False,
        'contact__solreffriction': False,
        'contact__type': False,
        'contact__worldid': False,
        'crb': True,
        'ctrl': True,
        'cvel': True,
        'efc__D': True,
        'efc__J': True,
        'efc__Ma': True,
        'efc__aref': True,
        'efc__force': True,
        'efc__frictionloss': True,
        'efc__id': True,
        'efc__margin': True,
        'efc__pos': True,
        'efc__state': True,
        'efc__type': True,
        'efc__vel': True,
        'energy': True,
        'eq_active': True,
        'flexedge_J': True,
        'flexedge_length': True,
        'flexedge_velocity': True,
        'flexvert_xpos': True,
        'geom_xmat': True,
        'geom_xpos': True,
        'light_xdir': True,
        'light_xpos': True,
        'mocap_pos': True,
        'mocap_quat': True,
        'naccdmax': False,
        'nacon': False,
        'naconmax': False,
        'ncollision': False,
        'ne': True,
        'nefc': True,
        'nf': True,
        'njmax': False,
        'nl': True,
        'nworld': False,
        'qLD': True,
        'qLDiagInv': True,
        'qM': True,
        'qacc': True,
        'qacc_smooth': True,
        'qacc_warmstart': True,
        'qfrc_actuator': True,
        'qfrc_applied': True,
        'qfrc_bias': True,
        'qfrc_constraint': True,
        'qfrc_damper': True,
        'qfrc_fluid': True,
        'qfrc_gravcomp': True,
        'qfrc_inverse': True,
        'qfrc_passive': True,
        'qfrc_smooth': True,
        'qfrc_spring': True,
        'qpos': True,
        'qvel': True,
        'sensordata': True,
        'site_xmat': True,
        'site_xpos': True,
        'solver_niter': True,
        'subtree_angmom': True,
        'subtree_com': True,
        'subtree_linvel': True,
        'ten_J': True,
        'ten_length': True,
        'ten_velocity': True,
        'ten_wrapadr': True,
        'ten_wrapnum': True,
        'time': True,
        'wrap_obj': True,
        'wrap_xpos': True,
        'xanchor': True,
        'xaxis': True,
        'xfrc_applied': True,
        'ximat': True,
        'xipos': True,
        'xmat': True,
        'xpos': True,
        'xquat': True,
    },
    'Model': {
        'M_colind': False,
        'M_rowadr': False,
        'M_rownnz': False,
        'actuator_acc0': False,
        'actuator_actadr': False,
        'actuator_actearly': False,
        'actuator_actlimited': False,
        'actuator_actnum': False,
        'actuator_actrange': True,
        'actuator_biasprm': True,
        'actuator_biastype': False,
        'actuator_cranklength': False,
        'actuator_ctrllimited': False,
        'actuator_ctrlrange': True,
        'actuator_dynprm': True,
        'actuator_dyntype': False,
        'actuator_forcelimited': False,
        'actuator_forcerange': True,
        'actuator_gainprm': True,
        'actuator_gaintype': False,
        'actuator_gear': True,
        'actuator_lengthrange': False,
        'actuator_trnid': False,
        'actuator_trntype': False,
        'actuator_trntype_body_adr': False,
        'block_dim__actuator_velocity': False,
        'block_dim__cholesky_factorize': False,
        'block_dim__cholesky_factorize_solve': False,
        'block_dim__cholesky_solve': False,
        'block_dim__contact_sort': False,
        'block_dim__energy_vel_kinetic': False,
        'block_dim__euler_dense': False,
        'block_dim__linesearch_iterative': False,
        'block_dim__qderiv_actuator_dense': False,
        'block_dim__ray': False,
        'block_dim__segmented_sort': False,
        'block_dim__tendon_velocity': False,
        'block_dim__update_gradient_JTDAJ_dense': False,
        'block_dim__update_gradient_JTDAJ_sparse': False,
        'block_dim__update_gradient_cholesky': False,
        'block_dim__update_gradient_cholesky_blocked': False,
        'body_branch_start': False,
        'body_branches': False,
        'body_conaffinity': False,
        'body_contype': False,
        'body_dofadr': False,
        'body_dofnum': False,
        'body_fluid_ellipsoid': False,
        'body_geomadr': False,
        'body_geomnum': False,
        'body_gravcomp': True,
        'body_inertia': True,
        'body_invweight0': True,
        'body_ipos': True,
        'body_iquat': True,
        'body_jntadr': False,
        'body_jntnum': False,
        'body_mass': True,
        'body_mocapid': False,
        'body_parentid': False,
        'body_pos': True,
        'body_quat': True,
        'body_rootid': False,
        'body_subtreemass': True,
        'body_tree': False,
        'body_treeid': False,
        'body_weldid': False,
        'cam_bodyid': False,
        'cam_fovy': True,
        'cam_intrinsic': True,
        'cam_mat0': True,
        'cam_mode': False,
        'cam_pos': True,
        'cam_pos0': True,
        'cam_poscom0': True,
        'cam_projection': False,
        'cam_quat': True,
        'cam_resolution': False,
        'cam_sensorsize': False,
        'cam_targetbodyid': False,
        'collision_sensor_adr': False,
        'dof_Madr': False,
        'dof_armature': True,
        'dof_bodyid': False,
        'dof_damping': True,
        'dof_frictionloss': True,
        'dof_invweight0': True,
        'dof_jntid': False,
        'dof_parentid': False,
        'dof_solimp': True,
        'dof_solref': True,
        'dof_treeid': False,
        'dof_tri_col': False,
        'dof_tri_row': False,
        'eq_active0': False,
        'eq_connect_adr': False,
        'eq_data': True,
        'eq_flex_adr': False,
        'eq_jnt_adr': False,
        'eq_obj1id': False,
        'eq_obj2id': False,
        'eq_objtype': False,
        'eq_solimp': True,
        'eq_solref': True,
        'eq_ten_adr': False,
        'eq_type': False,
        'eq_wld_adr': False,
        'exclude_signature': False,
        'flex_bending': False,
        'flex_damping': False,
        'flex_dim': False,
        'flex_edge': False,
        'flex_edgeadr': False,
        'flex_edgeflap': False,
        'flex_edgenum': False,
        'flex_elem': False,
        'flex_elemadr': False,
        'flex_elemedge': False,
        'flex_elemedgeadr': False,
        'flex_elemnum': False,
        'flex_stiffness': False,
        'flex_vertadr': False,
        'flex_vertbodyid': False,
        'flex_vertnum': False,
        'flexedge_J_colind': False,
        'flexedge_J_rowadr': False,
        'flexedge_J_rownnz': False,
        'flexedge_invweight0': False,
        'flexedge_length0': False,
        'geom_aabb': True,
        'geom_bodyid': False,
        'geom_conaffinity': False,
        'geom_condim': False,
        'geom_contype': False,
        'geom_dataid': False,
        'geom_fluid': False,
        'geom_friction': True,
        'geom_gap': True,
        'geom_group': False,
        'geom_margin': True,
        'geom_matid': True,
        'geom_pair_type_count': False,
        'geom_plugin_index': False,
        'geom_pos': True,
        'geom_priority': False,
        'geom_quat': True,
        'geom_rbound': True,
        'geom_rgba': True,
        'geom_size': True,
        'geom_solimp': True,
        'geom_solmix': True,
        'geom_solref': True,
        'geom_type': False,
        'has_fluid': False,
        'has_sdf_geom': False,
        'hfield_adr': False,
        'hfield_data': False,
        'hfield_ncol': False,
        'hfield_nrow': False,
        'hfield_size': False,
        'is_sparse': False,
        'jnt_actfrclimited': False,
        'jnt_actfrcrange': True,
        'jnt_actgravcomp': False,
        'jnt_axis': True,
        'jnt_bodyid': False,
        'jnt_dofadr': False,
        'jnt_limited': False,
        'jnt_limited_ball_adr': False,
        'jnt_limited_slide_hinge_adr': False,
        'jnt_margin': True,
        'jnt_pos': True,
        'jnt_qposadr': False,
        'jnt_range': True,
        'jnt_solimp': True,
        'jnt_solref': True,
        'jnt_stiffness': True,
        'jnt_type': False,
        'light_active': True,
        'light_bodyid': False,
        'light_castshadow': True,
        'light_dir': True,
        'light_dir0': True,
        'light_mode': False,
        'light_pos': True,
        'light_pos0': True,
        'light_poscom0': True,
        'light_targetbodyid': False,
        'light_type': True,
        'mapM2M': False,
        'mat_rgba': True,
        'mat_texid': True,
        'mat_texrepeat': True,
        'mesh_face': False,
        'mesh_faceadr': False,
        'mesh_graph': False,
        'mesh_graphadr': False,
        'mesh_normal': False,
        'mesh_normaladr': False,
        'mesh_polyadr': False,
        'mesh_polymap': False,
        'mesh_polymapadr': False,
        'mesh_polymapnum': False,
        'mesh_polynormal': False,
        'mesh_polynum': False,
        'mesh_polyvert': False,
        'mesh_polyvertadr': False,
        'mesh_polyvertnum': False,
        'mesh_quat': False,
        'mesh_vert': False,
        'mesh_vertadr': False,
        'mesh_vertnum': False,
        'mocap_bodyid': False,
        'nC': False,
        'nM': False,
        'na': False,
        'nacttrnbody': False,
        'nbody': False,
        'nbranch': False,
        'ncam': False,
        'neq': False,
        'nexclude': False,
        'nflex': False,
        'nflexedge': False,
        'nflexelem': False,
        'nflexelemdata': False,
        'nflexelemedge': False,
        'nflexvert': False,
        'ngeom': False,
        'ngravcomp': False,
        'nhfield': False,
        'nhfielddata': False,
        'njnt': False,
        'nlight': False,
        'nmat': False,
        'nmaxcondim': False,
        'nmaxmeshdeg': False,
        'nmaxpolygon': False,
        'nmaxpyramid': False,
        'nmesh': False,
        'nmeshface': False,
        'nmeshgraph': False,
        'nmeshnormal': False,
        'nmeshpoly': False,
        'nmeshpolymap': False,
        'nmeshpolyvert': False,
        'nmeshvert': False,
        'nmocap': False,
        'noct': False,
        'npair': False,
        'nplugin': False,
        'nq': False,
        'nrangefinder': False,
        'nsensor': False,
        'nsensorcollision': False,
        'nsensorcontact': False,
        'nsensordata': False,
        'nsensortaxel': False,
        'nsite': False,
        'ntendon': False,
        'ntree': False,
        'nu': False,
        'nv': False,
        'nv_pad': False,
        'nwrap': False,
        'nxn_geom_pair': False,
        'nxn_geom_pair_filtered': False,
        'nxn_pairid': False,
        'nxn_pairid_filtered': False,
        'oct_aabb': False,
        'oct_child': False,
        'oct_coeff': False,
        'opt__ccd_iterations': False,
        'opt__ccd_tolerance': True,
        'opt__cone': False,
        'opt__contact_sensor_maxmatch': False,
        'opt__density': True,
        'opt__disableflags': False,
        'opt__enableflags': False,
        'opt__graph_conditional': False,
        'opt__gravity': True,
        'opt__impratio_invsqrt': True,
        'opt__integrator': False,
        'opt__iterations': False,
        'opt__ls_iterations': False,
        'opt__ls_parallel': False,
        'opt__ls_parallel_min_step': False,
        'opt__ls_tolerance': True,
        'opt__magnetic': True,
        'opt__run_collision_detection': False,
        'opt__sdf_initpoints': False,
        'opt__sdf_iterations': False,
        'opt__solver': False,
        'opt__timestep': True,
        'opt__tolerance': True,
        'opt__viscosity': True,
        'opt__wind': True,
        'pair_dim': False,
        'pair_friction': True,
        'pair_gap': True,
        'pair_geom1': False,
        'pair_geom2': False,
        'pair_margin': True,
        'pair_solimp': True,
        'pair_solref': True,
        'pair_solreffriction': True,
        'plugin': False,
        'plugin_attr': False,
        'qLD_updates': False,
        'qM_fullm_i': False,
        'qM_fullm_j': False,
        'qM_mulm_col': False,
        'qM_mulm_madr': False,
        'qM_mulm_rowadr': False,
        'qM_tiles': False,
        'qpos0': True,
        'qpos_spring': True,
        'rangefinder_sensor_adr': False,
        'sensor_acc_adr': False,
        'sensor_adr': False,
        'sensor_adr_to_contact_adr': False,
        'sensor_collision_start_adr': False,
        'sensor_contact_adr': False,
        'sensor_cutoff': False,
        'sensor_datatype': False,
        'sensor_dim': False,
        'sensor_e_kinetic': False,
        'sensor_e_potential': False,
        'sensor_intprm': False,
        'sensor_limitfrc_adr': False,
        'sensor_limitpos_adr': False,
        'sensor_limitvel_adr': False,
        'sensor_objid': False,
        'sensor_objtype': False,
        'sensor_pos_adr': False,
        'sensor_rangefinder_adr': False,
        'sensor_rangefinder_bodyid': False,
        'sensor_refid': False,
        'sensor_reftype': False,
        'sensor_rne_postconstraint': False,
        'sensor_subtree_vel': False,
        'sensor_tendonactfrc_adr': False,
        'sensor_touch_adr': False,
        'sensor_type': False,
        'sensor_vel_adr': False,
        'site_bodyid': False,
        'site_pos': True,
        'site_quat': True,
        'site_size': False,
        'site_type': False,
        'stat__meaninertia': True,
        'taxel_sensorid': False,
        'taxel_vertadr': False,
        'ten_wrapadr_site': False,
        'ten_wrapnum_site': False,
        'tendon_actfrclimited': False,
        'tendon_actfrcrange': True,
        'tendon_adr': False,
        'tendon_armature': True,
        'tendon_damping': True,
        'tendon_frictionloss': True,
        'tendon_geom_adr': False,
        'tendon_invweight0': True,
        'tendon_jnt_adr': False,
        'tendon_length0': True,
        'tendon_lengthspring': True,
        'tendon_limited': False,
        'tendon_limited_adr': False,
        'tendon_margin': True,
        'tendon_num': False,
        'tendon_range': True,
        'tendon_site_pair_adr': False,
        'tendon_solimp_fri': True,
        'tendon_solimp_lim': True,
        'tendon_solref_fri': True,
        'tendon_solref_lim': True,
        'tendon_stiffness': True,
        'tree_bodynum': False,
        'tree_dofadr': False,
        'tree_dofnum': False,
        'wrap_geom_adr': False,
        'wrap_jnt_adr': False,
        'wrap_objid': False,
        'wrap_prm': False,
        'wrap_pulley_scale': False,
        'wrap_site_adr': False,
        'wrap_site_pair_adr': False,
        'wrap_type': False,
    },
    'Option': {
        'ccd_iterations': False,
        'ccd_tolerance': True,
        'cone': False,
        'contact_sensor_maxmatch': False,
        'density': True,
        'disableflags': False,
        'enableflags': False,
        'graph_conditional': False,
        'gravity': True,
        'impratio_invsqrt': True,
        'integrator': False,
        'iterations': False,
        'ls_iterations': False,
        'ls_parallel': False,
        'ls_parallel_min_step': False,
        'ls_tolerance': True,
        'magnetic': True,
        'run_collision_detection': False,
        'sdf_initpoints': False,
        'sdf_iterations': False,
        'solver': False,
        'timestep': True,
        'tolerance': True,
        'viscosity': True,
        'wind': True,
    },
    'Statistic': {'meaninertia': True},
}
