Allow Brax data structures to inherit from MJX data structures.

PiperOrigin-RevId: 601290735 Change-Id: I6dad1b6d7e6b3e71acecc36189ab59cc81ab90b4
2024-01-24 18:02:17 -08:00
parent e59605135a
commit a1e5ede266
1 changed files with 1 additions and 0 deletions
@@ -45,6 +45,7 @@ def dataclass(clz: _T) -> _T:
        field.type is jax.Array,
        dataclasses.is_dataclass(field.type),
        jax.Array in typing.get_args(field.type),
+        any(dataclasses.is_dataclass(a) for a in typing.get_args(field.type)),
    )):
      data_fields.append(field)
    else: