From d0f32a7fe6dab428b25de69ed04f79fcf326cf69 Mon Sep 17 00:00:00 2001
From: Baruch Tabanpour <btaba@google.com>
Date: Wed, 5 Nov 2025 13:31:57 -0800
Subject: [PATCH] Modify mjx.Model tex_data type.

PiperOrigin-RevId: 828606695
Change-Id: I708dcf8e75766bf16d339999a9e8c70c8f65fd0c
---
 doc/changelog.rst            | 8 ++++++++
 mjx/mujoco/mjx/_src/types.py | 2 +-
 2 files changed, 9 insertions(+), 1 deletion(-)

diff --git a/doc/changelog.rst b/doc/changelog.rst
index e26eab33..11ac1091 100644
--- a/doc/changelog.rst
+++ b/doc/changelog.rst
@@ -29,6 +29,14 @@ MJX
 
 - ``warp-lang`` optional dependency is updated to 1.10.0. ``pmap`` now works with MuJoCo Warp from MJX.
 
+.. admonition:: Breaking ABI changes
+   :class: attention
+
+   - ``mjx.Model.tex_data`` is now a numpy ndarray instead of a jax.Array, to avoid vmapping over this potentially
+     large array. This may break certain use-cases with Madrona MJX, but we are no longer supporting this codepath.
+     We will be migrating users to a Warp-based batch renderer.
+
+
 Version 3.3.7 (October 13, 2025)
 -----------------------------------
 
diff --git a/mjx/mujoco/mjx/_src/types.py b/mjx/mujoco/mjx/_src/types.py
index d1a25168..14004150 100644
--- a/mjx/mujoco/mjx/_src/types.py
+++ b/mjx/mujoco/mjx/_src/types.py
@@ -819,7 +819,7 @@ class Model(PyTreeNode):
   tex_width: np.ndarray
   tex_nchannel: np.ndarray
   tex_adr: np.ndarray
-  tex_data: jax.Array
+  tex_data: np.ndarray
   mat_rgba: jax.Array
   mat_texid: np.ndarray
   pair_dim: np.ndarray