Small changes for MJX interop. (google-deepmind#370)

* Small changes for MJX.

* Small changes for MJX.

Author: btaba

mujoco_warp/__init__.py

@@ -60,6 +60,7 @@
 from ._src.support import xfrc_accumulate as xfrc_accumulate
 from ._src.test_util import benchmark as benchmark
 from ._src.types import ConeType as ConeType
+from ._src.types import Constraint as Constraint
 from ._src.types import Contact as Contact
 from ._src.types import Data as Data
 from ._src.types import DisableBit as DisableBit

mujoco_warp/_src/types.py

@@ -28,6 +28,7 @@
 
 
 # TODO(team): add check that all wp.launch_tiled 'block_dim' settings are configurable
+@dataclasses.dataclass
 class BlockDim:
   """
   Block dimension 'block_dim' settings for wp.launch_tiled.