From a1e5ede266afa07db88acffe01e129e683056276 Mon Sep 17 00:00:00 2001 From: Baruch Tabanpour Date: Wed, 24 Jan 2024 18:02:17 -0800 Subject: [PATCH] Allow Brax data structures to inherit from MJX data structures. PiperOrigin-RevId: 601290735 Change-Id: I6dad1b6d7e6b3e71acecc36189ab59cc81ab90b4 --- mjx/mujoco/mjx/_src/dataclasses.py | 1 + 1 file changed, 1 insertion(+) diff --git a/mjx/mujoco/mjx/_src/dataclasses.py b/mjx/mujoco/mjx/_src/dataclasses.py index 5d95eb69..ea44d8f7 100644 --- a/mjx/mujoco/mjx/_src/dataclasses.py +++ b/mjx/mujoco/mjx/_src/dataclasses.py @@ -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: