[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

Author: danielsuo

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.