From 3e2d62c2484d699e38be243635d6f36210f8501b Mon Sep 17 00:00:00 2001
From: Erik Frey <erikfrey@google.com>
Date: Thu, 8 Feb 2024 15:00:47 -0800
Subject: [PATCH] Improve performance of getting and putting MJX device data
 with large numpy arrays.

ndarray.tobytes() is faster than converting arrays to tuples by orders of magnitude.

PiperOrigin-RevId: 605437697
Change-Id: I7d42c10e178b1aec80ff38a82e863e0701c15e62
---
 doc/changelog.rst                  |  8 ++++++++
 mjx/mujoco/mjx/_src/dataclasses.py | 12 +++++++-----
 2 files changed, 15 insertions(+), 5 deletions(-)

diff --git a/doc/changelog.rst b/doc/changelog.rst
index be62a872..8cbb2fb3 100644
--- a/doc/changelog.rst
+++ b/doc/changelog.rst
@@ -2,6 +2,14 @@
 Changelog
 =========
 
+Upcoming version (not yet released)
+-----------------------------------
+
+MJX
+^^^
+
+1. Improved performance of getting and putting device data by using a faster numpy array serialization method.
+
 Version 3.1.2 (February 05, 2024)
 -----------------------------------
 
diff --git a/mjx/mujoco/mjx/_src/dataclasses.py b/mjx/mujoco/mjx/_src/dataclasses.py
index ea44d8f7..1c71a9cd 100644
--- a/mjx/mujoco/mjx/_src/dataclasses.py
+++ b/mjx/mujoco/mjx/_src/dataclasses.py
@@ -58,12 +58,13 @@ def dataclass(clz: _T) -> _T:
   data_clz.replace = replace
 
   def iterate_clz_with_keys(x):
-    # numpy arrays are not hashable, so convert them to tuples for jit cache
-    to_tup = lambda x: tuple(x) if len(x.shape) == 1 else tuple(map(to_tup, x))
-
     def to_meta(field, obj):
       val = getattr(obj, field.name)
-      return (to_tup(val), val.dtype) if isinstance(val, np.ndarray) else val
+      # numpy arrays are not hashable so return raw bytes instead
+      if isinstance(val, np.ndarray):
+        return (val.tobytes(), val.dtype, val.shape)
+      else:
+        return val
 
     def to_data(field, obj):
       return (jax.tree_util.GetAttrKey(field.name), getattr(obj, field.name))
@@ -76,7 +77,8 @@ def dataclass(clz: _T) -> _T:
 
     def from_meta(field, meta):
       if field.type is np.ndarray:
-        return (field.name, np.array(meta[0], dtype=meta[1]))
+        arr = np.frombuffer(meta[0], dtype=meta[1]).reshape(meta[2])
+        return (field.name, arr)
       else:
         return (field.name, meta)