#mjx Make pytree registration optional.

This is a noop for MJX itself, but is useful for custom types extending MJX types.
This CL also ensures unused kwargs in `__init_subclass__` are passed to the parent class.

PiperOrigin-RevId: 728215373
Change-Id: I362d3fb8d4b7d679a3a00d96866d46a52d711369

mjx/mujoco/mjx/_src/dataclasses.py

@@ -34,7 +34,7 @@ def _jax_in_args(typ) -> bool:
   return False
 
 
-def dataclass(clz: _T) -> _T:
+def dataclass(clz: _T, register_as_pytree: bool) -> _T:
   """Wraps a dataclass with metadata for which fields are pytrees.
 
   This is based off flax.struct.dataclass, but instead of using field
@@ -48,6 +48,9 @@ def dataclass(clz: _T) -> _T:
     the resulting dataclass, registered with Jax
   """
   data_clz = dataclasses.dataclass(frozen=True)(clz)
+  data_clz.replace = dataclasses.replace
+
+  if register_as_pytree:
     meta_fields, data_fields = [], []
     for field in dataclasses.fields(data_clz):
       if _jax_in_args(field.type):
@@ -55,12 +58,6 @@ def dataclass(clz: _T) -> _T:
       else:
         meta_fields.append(field)
 
-  def replace(self, **updates):
-    """Returns a new object replacing the specified fields with new values."""
-    return dataclasses.replace(self, **updates)
-
-  data_clz.replace = replace
-
     def iterate_clz_with_keys(x):
       def to_meta(field, obj):
         val = getattr(obj, field.name)
@@ -109,8 +106,9 @@ class PyTreeNode:
   This base class additionally avoids type checking errors when using PyType.
   """
 
-  def __init_subclass__(cls):
-    dataclass(cls)
+  def __init_subclass__(cls, register_as_pytree: bool = True, **kwargs):
+    super().__init_subclass__(**kwargs)
+    dataclass(cls, register_as_pytree=register_as_pytree)
 
   def __init__(self, *args, **kwargs):
     # stub for pytype