From 31022d6cd7dd525ed11a04d84cd3936228499173 Mon Sep 17 00:00:00 2001
From: Daniel Suo <dsuo@google.com>
Date: Fri, 26 Sep 2025 07:36:11 -0700
Subject: [PATCH] [trax] Explicitly set `jax_pmap_shmap_merge=False`.

`trainer._multi_device_update_fn` uses `jax.pmap` and when `jax_pmap_shmap_merge=True`, `jax.pmap` requires inputs be explicitly sharded as the underlying `jax.jit` expects.

This would need to be fixed if `jax_pmap_shmap_merge=True`.

PiperOrigin-RevId: 811810947
---
 trax/optimizers/trainer.py | 6 ++++++
 1 file changed, 6 insertions(+)

diff --git a/trax/optimizers/trainer.py b/trax/optimizers/trainer.py
index 4f5a2d5e3..c0d9a069a 100644
--- a/trax/optimizers/trainer.py
+++ b/trax/optimizers/trainer.py
@@ -32,6 +32,12 @@
 from trax.layers import combinators as cb
 
 
+# NOTE(dsuo): _multi_device_update_fn is not compatible with
+# jax_pmap_shmap_merge=True because `jax.pmap` requires inputs to be explicitly
+# sharded as the underlying `jax.jit` expects.
+jax.config.update('jax_pmap_shmap_merge', False)
+
+
 class Trainer:
   """Multi-device accelerated trainer.