From d561995c037a8b54a7861556684a05acdedcdf71 Mon Sep 17 00:00:00 2001
From: Erik Frey <erikfrey@google.com>
Date: Wed, 14 Feb 2024 13:57:03 -0800
Subject: [PATCH] Use less brittle search to determine whether an MJX field is
 a data or meta pytree node.  Use correct type annotation for geom_convex_*
 fields in mjx.Model.

PiperOrigin-RevId: 607096855
Change-Id: I2e14b9b63e676b4ac3ec2a12aa07316e1d019a79
---
 mjx/mujoco/mjx/_src/dataclasses.py | 19 ++++++++++++-------
 mjx/mujoco/mjx/_src/types.py       | 10 +++++-----
 2 files changed, 17 insertions(+), 12 deletions(-)

diff --git a/mjx/mujoco/mjx/_src/dataclasses.py b/mjx/mujoco/mjx/_src/dataclasses.py
index 1c71a9cd..5a80ab4c 100644
--- a/mjx/mujoco/mjx/_src/dataclasses.py
+++ b/mjx/mujoco/mjx/_src/dataclasses.py
@@ -18,13 +18,23 @@ import copy
 import dataclasses
 
 import typing
-from typing import Any, Dict, Optional, Sequence, TypeVar
+from typing import Any, Dict, Optional, Sequence, TypeVar, Union
 import jax
 import numpy as np
 
 _T = TypeVar('_T')
 
 
+def _jax_in_args(typ) -> bool:
+  if typ is jax.Array:
+    return True
+  if dataclasses.is_dataclass(typ):
+    return any(_jax_in_args(f.type) for f in dataclasses.fields(typ))
+  if typing.get_origin(typ) in (list, dict, Union, set):
+    return any(_jax_in_args(t) for t in typing.get_args(typ))
+  return False
+
+
 def dataclass(clz: _T) -> _T:
   """Wraps a dataclass with metadata for which fields are pytrees.
 
@@ -41,12 +51,7 @@ def dataclass(clz: _T) -> _T:
   data_clz = dataclasses.dataclass(frozen=True)(clz)
   meta_fields, data_fields = [], []
   for field in dataclasses.fields(data_clz):
-    if any((
-        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)),
-    )):
+    if _jax_in_args(field.type):
       data_fields.append(field)
     else:
       meta_fields.append(field)
diff --git a/mjx/mujoco/mjx/_src/types.py b/mjx/mujoco/mjx/_src/types.py
index bb04c35e..69878e67 100644
--- a/mjx/mujoco/mjx/_src/types.py
+++ b/mjx/mujoco/mjx/_src/types.py
@@ -15,7 +15,7 @@
 """Base types used in MJX."""
 
 import enum
-from typing import Sequence
+from typing import List, Optional
 
 import jax
 import jax.numpy as jp
@@ -498,10 +498,10 @@ class Model(PyTreeNode):
   pair_dim: np.ndarray
   pair_geom1: np.ndarray
   pair_geom2: np.ndarray
-  geom_convex_face: Sequence[jax.Array]
-  geom_convex_vert: Sequence[jax.Array]
-  geom_convex_edge: Sequence[jax.Array]
-  geom_convex_facenormal: Sequence[jax.Array]
+  geom_convex_face: List[Optional[jax.Array]]
+  geom_convex_vert: List[Optional[jax.Array]]
+  geom_convex_edge: List[Optional[jax.Array]]
+  geom_convex_facenormal: List[Optional[jax.Array]]
   pair_solref: jax.Array
   pair_solreffriction: jax.Array
   pair_solimp: jax.Array