No public description

PiperOrigin-RevId: 877435039

Author: danielsuo

examples/wmt/train.py

@@ -363,10 +363,9 @@ def pre_pmap(xs):
   def post_pmap(xs):
     # Avoid degraded performance under the new jax.pmap. See
     # https://docs.jax.dev/en/latest/migrate_pmap.html#int-indexing-into-sharded-arrays.
-    if jax.config.jax_pmap_shmap_merge:
-      return jax.tree_util.tree_map(
-          lambda x: x.addressable_shards[0].data.squeeze(0), xs)
-    return jax.tree_util.tree_map(lambda x: x[0], xs)
+    return jax.tree_util.tree_map(
+        lambda x: x.addressable_shards[0].data.squeeze(0), xs
+    )
 
   return post_pmap(host_psum(pre_pmap(in_tree)))
 

flax/jax_utils.py

@@ -47,22 +47,22 @@ def replicate(tree, devices=None):
 
 def unreplicate(tree):
   """Returns a single instance of a replicated array."""
-  if jax.config.jax_pmap_shmap_merge:
-    def _unreplicate_one(x):
-      # Avoid degraded performance under the new jax.pmap. See
-      # https://docs.jax.dev/en/latest/migrate_pmap.html#int-indexing-into-sharded-arrays.
-      # Handle 0-dimensional (scalar) arrays - cannot index into them
-      if hasattr(x, 'ndim') and x.ndim == 0:
-        return x
-      if (not hasattr(x, 'sharding') or
-          isinstance(x.sharding, jax.sharding.SingleDeviceSharding) or
-          len(jax.local_devices()) == 1):
-        return x[0]
-      if x.sharding.is_fully_replicated:
-        return x.addressable_shards[0].data
-      return x.addressable_shards[0].data.squeeze(0)
-    return jax.tree_util.tree_map(_unreplicate_one, tree)
-  return jax.tree_util.tree_map(lambda x: x[0], tree)
+  def _unreplicate_one(x):
+    # Avoid degraded performance under the new jax.pmap.
+    # Handle 0-dimensional (scalar) arrays - cannot index into them
+    if hasattr(x, 'ndim') and x.ndim == 0:
+      return x
+    if (
+        not hasattr(x, 'sharding')
+        or isinstance(x.sharding, jax.sharding.SingleDeviceSharding)
+        or len(jax.local_devices()) == 1
+    ):
+      return x[0]
+    if x.sharding.is_fully_replicated:
+      return x.addressable_shards[0].data
+    return x.addressable_shards[0].data.squeeze(0)
+
+  return jax.tree_util.tree_map(_unreplicate_one, tree)
 
 
 def pmean(xs, axis_name):

flax/training/common_utils.py

@@ -80,13 +80,9 @@ def get_metrics(device_metrics):
   """
   # We select the first element of x in order to get a single copy of a
   # device-replicated metric.
-  # Avoid degraded performance under the new jax.pmap. See
-  # https://docs.jax.dev/en/latest/migrate_pmap.html#int-indexing-into-sharded-arrays.
-  if jax.config.jax_pmap_shmap_merge:
-    device_metrics = jax.tree_util.tree_map(
-        lambda x: x.addressable_shards[0].data.squeeze(0), device_metrics)
-  else:
-    device_metrics = jax.tree_util.tree_map(lambda x: x[0], device_metrics)
+  device_metrics = jax.tree_util.tree_map(
+      lambda x: x.addressable_shards[0].data.squeeze(0), device_metrics
+  )
   metrics_np = jax.device_get(device_metrics)
   return stack_forest(metrics_np)