From 86bca81d18cc4612615b42e079694cbe5632a3e4 Mon Sep 17 00:00:00 2001
From: Baruch Tabanpour <btaba@google.com>
Date: Thu, 19 Feb 2026 13:17:23 -0800
Subject: [PATCH] Update mjx docs and changelog.

PiperOrigin-RevId: 872535000
Change-Id: I51b05a22ac553c34b436de63f43c7efca6607724
---
 doc/changelog.rst |  9 +++++++++
 doc/mjx.rst       | 29 ++++++++++++++++++++++++++---
 2 files changed, 35 insertions(+), 3 deletions(-)

diff --git a/doc/changelog.rst b/doc/changelog.rst
index 79a1929e..eeb6fab0 100644
--- a/doc/changelog.rst
+++ b/doc/changelog.rst
@@ -2,6 +2,15 @@
 Changelog
 =========
 
+Upcoming version (not yet released)
+-----------------------------------
+
+MJX
+^^^
+
+- Add batch rendering support for MJX-Warp. See the :ref:`MJX-Warp batch rendering<MjxWarpBatchRendering>` section for details.
+
+
 Version 3.5.0 (February 12, 2026)
 ---------------------------------
 
diff --git a/doc/mjx.rst b/doc/mjx.rst
index 8825d869..a68d57b1 100644
--- a/doc/mjx.rst
+++ b/doc/mjx.rst
@@ -207,6 +207,8 @@ excessive graph captures in the JAX-Warp FFI layer.
      - 0.65M
 
 
+.. _MjxWarpBatchRendering:
+
 MJX-Warp Batch Rendering
 ~~~~~~~~~~~~~~~~~~~~~~~~
 
@@ -218,9 +220,9 @@ Note that the number of parallel worlds (``nworld``) is fixed when creating the
 
 .. code-block:: python
 
-    from mujoco.mjx import io
+    from mujoco.mjx import create_render_context
 
-    rc = io.create_render_context(
+    rc = create_render_context(
         mjm=m,
         nworld=nworld,
         cam_res=(width, height),
@@ -247,12 +249,33 @@ volume hierarchy (BVH) and executing the raycaster:
         pixels, _ = mjx.render(mx, d, rc)
 
         # 3. Extract the RGB tensor for the first camera (index 0)
-        rgb = get_rgb(rc, pixels, 0)
+        rgb = get_rgb(rc, 0, pixels)
 
         # CAVEAT: Always return or use the updated `d` in your computation graph.
         # Otherwise, JAX's dead-code elimination will optimize away the refit_bvh call!
         return rgb, d
 
+Multi-GPU rendering with ``pmap``
+^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+
+To render across multiple GPUs, create a render context **per device** by passing ``devices`` to
+:func:`create_render_context <mujoco.mjx.create_render_context>`.
+
+.. code-block:: python
+
+    ndevices = jax.local_device_count()
+    nworld_per_device = nworld // ndevices
+
+    # Create one render context for all devices
+    rc = create_render_context(
+        mjm=m,
+        nworld=nworld_per_device,
+        devices=[f'cuda:{i}' for i in range(ndevices)],
+        cam_res=(width, height),
+    )
+
+Then use ``jax.pmap`` to parallelize the rendering across devices. See the complete example in
+`visualize_render.py <https://github.com/google-deepmind/mujoco/blob/main/mjx/mujoco/mjx/warp/visualize_render.py>`__.
 
 .. _MjxJAX: