Merge pull request #3295 from tkelestemur:tarik/fix-mjx-codegen-python311

PiperOrigin-RevId: 922794565
Change-Id: Ibd42328d739d3882e5c418cf69fc9787572cdb4a

mjx/mujoco/mjx/codegen/generate_warp_shim.py

@@ -215,7 +215,7 @@ def _warp_function(
       info = mjwarp_field_info[f]
       expected_type = _clean_type(info.expected_type)
       fn_args_model.append((f'{f}: {expected_type},', info.param_order))
-      fn_assignments.append(f'  _m.{f.replace('__', '.')} = {f}')
+      fn_assignments.append(f"  _m.{f.replace('__', '.')} = {f}")
     fn_args_model = sorted(fn_args_model, key=lambda x: x[1])
     fn_args_model = ['# Model'] + [f[0] for f in fn_args_model]
 
@@ -233,7 +233,7 @@ def _warp_function(
       param_order = info.param_order
       expected_type = _clean_type(info.expected_type)
       fn_args_data.append((f'{f}: {expected_type},', (is_out, param_order)))
-      fn_assignments.append(f'  _d.{f.replace('__', '.')} = {f}')
+      fn_assignments.append(f"  _d.{f.replace('__', '.')} = {f}")
     fn_args_data = sorted(fn_args_data, key=lambda x: x[1])
     fn_args_data = ['# Data'] + [f[0] for f in fn_args_data]
 
@@ -418,6 +418,7 @@ def create_jax_warp_shim(
   fn_args_raw_str = '\n'.join(['    ' + arg for arg in fn_args_raw])
   warp_fn_args = [arg.split(':')[0] for arg in fn_args_raw if '#' not in arg]  # pytype: disable=attribute-error
 
+  fn_assignments_str = '\n'.join(fn_assignments)
   src += f"""
 @ffi.format_args_for_warp
 def _{fn_name}_shim(
@@ -428,7 +429,7 @@ def _{fn_name}_shim(
   _m.callback = _cb
   _d.efc = _e
   _d.contact = _c
-{'\n'.join(fn_assignments)}
+{fn_assignments_str}
   {fn_call}
   """
   src += '\n\n'

mjx/mujoco/mjx/codegen/generate_warp_types.py

@@ -67,7 +67,8 @@ def _to_py_string(value, indent=0):
   next_indent_str = '  ' * (indent + 1)
   if isinstance(value, tuple):
     items = [_to_py_string(item, indent) for item in value]
-    return f'({', '.join(items)})'
+    joined_items = ', '.join(items)
+    return f'({joined_items})'
 
   if isinstance(value, type):
     if value.__module__ == 'builtins':
@@ -79,12 +80,14 @@ def _to_py_string(value, indent=0):
         f'\n{next_indent_str}{repr(k)}: {_to_py_string(v, indent + 1)}'
         for k, v in sorted(value.items(), key=lambda x: x[0])
     ]
-    return f'{{{','.join(items)}\n{indent_str}}}'
+    joined_items = ','.join(items)
+    return f'{{{joined_items}\n{indent_str}}}'
 
   if isinstance(value, set):
     items = sorted([_to_py_string(item, indent) for item in value])
     items = [f'\n{next_indent_str}{item}' for item in items]
-    return f'{{{",".join(items)}\n{indent_str}}}'
+    joined_items = ','.join(items)
+    return f'{{{joined_items}\n{indent_str}}}'
 
   return repr(value)
 

mjx/mujoco/mjx/warp/bvh.py

@@ -14,17 +14,19 @@
 # ==============================================================================
 
 """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
 from mujoco.mjx.warp.render_context import _MJX_RENDER_CONTEXT_BUFFERS
 from mujoco.mjx.warp.render_context import RenderContextPytree
-import mujoco.mjx.third_party.mujoco_warp as mjwarp
-from mujoco.mjx.third_party.mujoco_warp._src import types as mjwp_types
-import warp as wp
-
 
 _m = mjwarp.Model(
     **{f.name: None for f in dataclasses.fields(mjwarp.Model) if f.init}
@@ -48,6 +50,7 @@ _cb = mjwp_types.Callback(
     **{f.name: None for f in dataclasses.fields(mjwp_types.Callback) if f.init}
 )
 
+
 @ffi.format_args_for_warp
 def _refit_bvh_shim(
     # Model

mjx/mujoco/mjx/warp/collision_driver.py

@@ -14,15 +14,17 @@
 # ==============================================================================
 
 """DO NOT EDIT. This file is auto-generated."""
+
 import dataclasses
 import functools
+
 import jax
-from mujoco.mjx._src import types
-from mujoco.mjx.warp import ffi
-import mujoco.mjx.third_party.mujoco_warp as mjwarp
-from mujoco.mjx.third_party.mujoco_warp._src import types as mjwp_types
 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}
@@ -46,6 +48,7 @@ _cb = mjwp_types.Callback(
     **{f.name: None for f in dataclasses.fields(mjwp_types.Callback) if f.init}
 )
 
+
 @ffi.format_args_for_warp
 def _collision_shim(
     # Model

mjx/mujoco/mjx/warp/forward.py

@@ -14,15 +14,17 @@
 # ==============================================================================
 
 """DO NOT EDIT. This file is auto-generated."""
+
 import dataclasses
 import functools
+
 import jax
-from mujoco.mjx._src import types
-from mujoco.mjx.warp import ffi
-import mujoco.mjx.third_party.mujoco_warp as mjwarp
-from mujoco.mjx.third_party.mujoco_warp._src import types as mjwp_types
 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}

mjx/mujoco/mjx/warp/render.py

@@ -14,17 +14,19 @@
 # ==============================================================================
 
 """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
 from mujoco.mjx.warp.render_context import _MJX_RENDER_CONTEXT_BUFFERS
 from mujoco.mjx.warp.render_context import RenderContextPytree
-import mujoco.mjx.third_party.mujoco_warp as mjwarp
-from mujoco.mjx.third_party.mujoco_warp._src import types as mjwp_types
-import warp as wp
-
 
 _m = mjwarp.Model(
     **{f.name: None for f in dataclasses.fields(mjwarp.Model) if f.init}
@@ -48,6 +50,7 @@ _cb = mjwp_types.Callback(
     **{f.name: None for f in dataclasses.fields(mjwp_types.Callback) if f.init}
 )
 
+
 @ffi.format_args_for_warp
 def _render_shim(
     # Model

mjx/mujoco/mjx/warp/smooth.py

@@ -14,15 +14,17 @@
 # ==============================================================================
 
 """DO NOT EDIT. This file is auto-generated."""
+
 import dataclasses
 import functools
+
 import jax
-from mujoco.mjx._src import types
-from mujoco.mjx.warp import ffi
-import mujoco.mjx.third_party.mujoco_warp as mjwarp
-from mujoco.mjx.third_party.mujoco_warp._src import types as mjwp_types
 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}
@@ -46,6 +48,7 @@ _cb = mjwp_types.Callback(
     **{f.name: None for f in dataclasses.fields(mjwp_types.Callback) if f.init}
 )
 
+
 @ffi.format_args_for_warp
 def _kinematics_shim(
     # Model

mjx/mujoco/mjx/warp/types.py

@@ -15,14 +15,18 @@
 """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
+
+from mujoco.mjx._src import dataclasses as mjx_dataclasses
+
 if typing.TYPE_CHECKING:
   GraphMode = int
 
@@ -33,6 +37,7 @@ if typing.TYPE_CHECKING:
 else:
   try:
     from warp._src.jax_experimental.ffi import GraphMode
+
     from mujoco.mjx.third_party.mujoco_warp._src import types as mjwp_types
 
     Callback = mjwp_types.Callback
@@ -41,6 +46,7 @@ else:
     Callback = None
 PyTreeNode = mjx_dataclasses.PyTreeNode
 
+
 @dataclasses.dataclass(frozen=True)
 @tree_util.register_pytree_node_class
 class TileSet:
@@ -52,6 +58,7 @@ class TileSet:
     adr: address of each tile in the set
     size: size of all the tiles in this set
   """
+
   adr: np.ndarray
   size: int
 
@@ -72,6 +79,7 @@ class BlockDim:
 
   TODO(team): experimental and may be removed
   """
+
   actuator_velocity: int
   cholesky_factorize: int
   cholesky_factorize_solve: int
@@ -102,10 +110,13 @@ class BlockDim:
 
 class StatisticWarp(PyTreeNode):
   """Derived fields from Statistic."""
+
   meaninertia: jax.Array
 
+
 class OptionWarp(PyTreeNode):
   """Derived fields from Option."""
+
   broadphase: int
   broadphase_filter: int
   ccd_iterations: int
@@ -120,8 +131,10 @@ class OptionWarp(PyTreeNode):
   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
@@ -283,8 +296,10 @@ class ModelWarp(PyTreeNode):
   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
@@ -363,6 +378,8 @@ class DataWarp(PyTreeNode):
   wrap_obj: jax.Array
   wrap_xpos: jax.Array
   shape = property(lambda self: self.cacc.shape)
+
+
 DATA_NON_VMAP = {
     'contact__dim',
     'contact__dist',
@@ -390,6 +407,7 @@ DATA_NON_VMAP = {
     'nworld',
 }
 
+
 def _to_elt(cont, _, d, axis):
   return DataWarp(**{
       f.name: (