MJX: Use content-hash for np.ndarray pytree metadata fields

Previously mjx dataclasses would return np.ndarray as bytes for jax tracing to hash since they are not inherently hashable. This meant that any tracing that was cached would hold on to a copy of the numpy arrays in the dataclass. Children such as mjx.Model that store large numpy arrays would end up duplicating that data 3+ times in some cases.

This eliminates O(N * array_size) memory duplication across N cached
pytree traces, saving a lot of memory on models with heavy mesh/texture data.

PiperOrigin-RevId: 880985898
Change-Id: I58d4e91fda4112e4c633818f92907d20138e962a

mjx/mujoco/mjx/_src/dataclasses.py

@@ -16,15 +16,42 @@
 
 import copy
 import dataclasses
+import hashlib
 import typing
 from typing import Dict, Optional, Sequence, Tuple, TypeVar, Union
-import warnings
 import jax
 import numpy as np
 
 _T = TypeVar('_T')
 
 
+class _NumPyArrayHashWrapper:
+  """A wrapper for NumPy arrays to make them hashable based on content.
+
+  This class is used to allow NumPy arrays to be part of the metadata in a Jax
+  PyTree registration, as metadata must be hashable. The hash is based on the
+  array's content, dtype, and shape.
+  """
+  __slots__ = ('_hash_key', 'array')
+
+  def __init__(self, arr: np.ndarray):
+    if arr.size == 0:
+      h = hashlib.sha256(b'').hexdigest()
+    else:
+      contiguous = np.ascontiguousarray(arr)
+      h = hashlib.sha256(contiguous.data.cast('B')).hexdigest()
+    self._hash_key = (h, arr.dtype, arr.shape)
+    self.array = arr
+
+  def __hash__(self):
+    return hash(self._hash_key)
+
+  def __eq__(self, other):
+    if not isinstance(other, _NumPyArrayHashWrapper):
+      return NotImplemented
+    return self._hash_key == other._hash_key
+
+
 def _jax_in_args(typ) -> bool:
   if typ is jax.Array:
     return True
@@ -63,17 +90,15 @@ def dataclass(clz: _T, register_as_pytree: bool) -> _T:
       def to_meta(field, obj):
         val = getattr(obj, field.name)
         if isinstance(val, np.ndarray):
-          # numpy arrays are not hashable so return raw bytes instead
-          return (val.tobytes(), val.dtype, val.shape)
+          return _NumPyArrayHashWrapper(val)
         if typing.get_origin(field.type) == tuple:
-          # variadic tuples of numpy arrays
           type_args = typing.get_args(field.type)
           if (
               len(type_args) == 2
               and type_args[0] == np.ndarray
               and type_args[1] == ...
           ):
-            return tuple((v.tobytes(), v.dtype, v.shape) for v in val)
+            return tuple(_NumPyArrayHashWrapper(v) for v in val)
         return val
 
       def to_data(field, obj):
@@ -87,8 +112,7 @@ def dataclass(clz: _T, register_as_pytree: bool) -> _T:
 
       def from_meta(field, meta):
         if field.type is np.ndarray:
-          arr = np.frombuffer(meta[0], dtype=meta[1]).reshape(meta[2])
-          return (field.name, arr)
+          return (field.name, meta.array)
         if typing.get_origin(field.type) == tuple:
           type_args = typing.get_args(field.type)
           if (
@@ -98,9 +122,7 @@ def dataclass(clz: _T, register_as_pytree: bool) -> _T:
           ):
             return (
                 field.name,
-                tuple(
-                    np.frombuffer(m[0], dtype=m[1]).reshape(m[2]) for m in meta
-                ),
+                tuple(m.array for m in meta),
             )
         return (field.name, meta)
 

mjx/mujoco/mjx/_src/dataclasses_test.py

@@ -30,6 +30,12 @@ class Obj(dataclasses.PyTreeNode):
   f: tuple[jax.Array, ...]
 
 
+class LargeArrayNode(dataclasses.PyTreeNode):
+  array_a: np.ndarray
+  array_b: np.ndarray
+  array_c: np.ndarray
+
+
 class DataclassesTest(absltest.TestCase):
 
   def test_pytree_structure(self):
@@ -64,6 +70,48 @@ class DataclassesTest(absltest.TestCase):
     # ensure hashable meta
     hash(meta)
 
+  def test_metadata_equality(self):
+    array_size = 1_000
+    data = np.random.rand(array_size, 3).astype(np.float32)
+    tex = np.random.randint(0, 255, (1024, 1024, 3), dtype=np.uint8)
+    obj1 = LargeArrayNode(
+        array_a=data.copy(),
+        array_b=data.copy(),
+        array_c=tex.copy(),
+    )
+    obj2 = LargeArrayNode(
+        array_a=data.copy(),
+        array_b=data.copy(),
+        array_c=tex.copy(),
+    )
+
+    _, meta1 = jax.tree_util.tree_flatten(obj1)
+    _, meta2 = jax.tree_util.tree_flatten(obj2)
+
+    self.assertEqual(
+        meta1,
+        meta2,
+        'Two objects with identical numpy data should produce equal pytree'
+        ' metadata (same trace cache key).',
+    )
+
+  def test_metadata_does_not_copy_arrays(self):
+    obj = LargeArrayNode(
+        array_a=np.random.rand(100, 3).astype(np.float32),
+        array_b=np.random.rand(100, 3).astype(np.float32),
+        array_c=np.random.randint(0, 255, (16, 16, 3), dtype=np.uint8),
+    )
+
+    leaves, meta = jax.tree_util.tree_flatten(obj)
+    reconstructed = meta.unflatten(leaves)
+
+    self.assertIs(
+        reconstructed.array_a,
+        obj.array_a,
+        'Metadata should reference the original array.',
+    )
+    self.assertIs(reconstructed.array_c, obj.array_c)
+
 
 if __name__ == '__main__':
   absltest.main()