fix mjx codegen and rerun

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)