Implementation:Google deepmind Mujoco MJX Warp Types

Knowledge Sources	Google_deepmind_Mujoco
Domains	Type System, Data Structures, JAX-Warp Bridge, GPU Computing
Last Updated	2026-02-15 04:00 GMT

Overview

Auto-generated type definitions module that declares JAX-compatible PyTree dataclasses for MJX Warp bridge data structures including tiling configurations and block dimension settings.

Description

types.py defines the core data structures used throughout the MJX Warp bridge layer. It includes TileSet for describing decomposable block-diagonal matrix tiling configurations (with address arrays and tile sizes), BlockDim for specifying GPU kernel launch block dimensions across all tiled Warp kernels (actuator_velocity, cholesky_factorize, euler_dense, etc.), and imports GraphMode from the Warp JAX experimental FFI module. All dataclasses are registered as JAX PyTree nodes using @tree_util.register_pytree_node_class to enable seamless integration with JAX transformations (jit, vmap, grad).

Usage

These types are used by all MJX Warp bridge modules (forward, smooth, collision_driver) as parameter types in shim function signatures. The BlockDim dataclass controls GPU thread block sizes for each kernel, and TileSet describes sparse matrix decomposition layouts for efficient GPU computation.

Code Reference

Source Location

• Repository: Google_deepmind_Mujoco
• File: mjx/mujoco/mjx/warp/types.py
• Lines: 1-1467

Key Functions

@dataclasses.dataclass(frozen=True)
@tree_util.register_pytree_node_class
class TileSet:
    """Tiling configuration for decomposable block diagonal matrix."""
    adr: np.ndarray   # address of each tile in the set
    size: int          # size of all the tiles in this set

@dataclasses.dataclass(frozen=True)
@tree_util.register_pytree_node_class
class BlockDim:
    """Block dimension 'block_dim' settings for wp.launch_tiled."""
    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
    ...

Import

from mujoco.mjx.warp import types
from mujoco.mjx.warp.types import TileSet, BlockDim, GraphMode

I/O Contract

Inputs

Name	Type	Required	Description
adr	np.ndarray	Yes	Tile address array for TileSet construction
size	int	Yes	Uniform tile size for TileSet construction
(BlockDim fields)	int	Yes	Per-kernel block dimension values for GPU thread configuration

Outputs

Name	Type	Description
TileSet	frozen dataclass / PyTreeNode	Immutable tiling configuration registered as a JAX PyTree
BlockDim	frozen dataclass / PyTreeNode	Immutable block dimension settings registered as a JAX PyTree
GraphMode	IntEnum	Graph capture mode (NONE, JAX, WARP, WARP_STAGED, WARP_STAGED_EX)

Related Pages

• Google_deepmind_Mujoco_MJX_Warp_Forward - Uses BlockDim and TileSet in forward dynamics shim
• Google_deepmind_Mujoco_MJX_Warp_Smooth - Uses BlockDim in kinematics shim
• Google_deepmind_Mujoco_MJX_Warp_Collision_Driver - Uses BlockDim in collision shim
• Google_deepmind_Mujoco_JAX_Experimental_FFI - Source of GraphMode enum