[kfac_jax] Prepare for jax_pmap_shmap_merge=True.

Condition on `jax_pmap_shmap_merge` to grab explicitly grab the first shard rather than rely on `x[0]`. See https://docs.jax.dev/en/latest/migrate_pmap.html#int-indexing-into-sharded-arrays for more information.

NOTE: `kfac_jax.utils.index_if_not_scalar` has been removed and inlined into `kfac_jax.utils.get_first`.
- If you need to grab the first shard of a semantically replicated array (i.e., each shard is really a replica), use `kfac_jax.utils.get_first`.
- If you really are trying to index into an array, use the usual array indexing operators (i.e., `x[0]`).
PiperOrigin-RevId: 846329906

Author: danielsuo

examples/training.py

@@ -520,7 +520,8 @@ def format_path_entry(entry: Any) -> str:
     logging.info("%s %s %s", "=" * 20, "Parameters", "=" * 20)
     for path, var in jax.tree_util.tree_flatten_with_path(self._params)[0]:
       # Because of pmap
-      var = var[0]
+      var = kfac_jax.utils.get_first(var)
+      assert isinstance(var, Array)  # params are always arrays
       logging.info(
           "%s - %s, %s",
           "-".join(format_path_entry(p) for p in path),
@@ -541,7 +542,8 @@ def format_path_entry(entry: Any) -> str:
         # For __class__ entries
         continue
       # Because of pmap
-      var = var[0]
+      var = kfac_jax.utils.get_first(var)
+      assert isinstance(var, Array)  # optimizer state entries are always arrays
       logging.info(
           "%s - %s, %s",
           "/".join(format_path_entry(p) for p in path),
@@ -671,7 +673,10 @@ def train_step(self, global_step: Array, rng: PRNGKey) -> dict[str, Numeric]:
       for i in range(gathered_stat.shape[0]):
         stats[f"{name}_{i}"] = jnp.array([gathered_stat[i]])
 
-    stats = jax.tree_util.tree_map(functools.partial(jnp.mean, axis=0), stats)
+    if jax.config.jax_pmap_shmap_merge:
+      stats = kfac_jax.utils.get_first(stats)
+    else:
+      stats = jax.tree_util.tree_map(functools.partial(jnp.mean, axis=0), stats)
 
     self._python_step += 1
     stats["progress"] = self.progress(self._python_step)

kfac_jax/_src/utils/__init__.py

@@ -73,7 +73,6 @@
 psum_if_pmap = parallel.psum_if_pmap
 pmap_mean = parallel.pmap_mean
 pmap_sum = parallel.pmap_sum
-index_if_not_scalar = parallel.index_if_not_scalar
 get_first = parallel.get_first
 get_mean = parallel.get_mean
 get_sum = parallel.get_sum

kfac_jax/_src/utils/parallel.py

@@ -82,24 +82,27 @@ def psum_if_pmap(obj: TArrayTree, axis_name: str | None) -> TArrayTree:
 pmap_sum = jax.pmap(lambda x: lax.psum(x, "i"), axis_name="i")
 
 
-def index_if_not_scalar(value: Numeric, index: int = 0) -> Numeric:
-  """Index `value` at axis 0 if it is not a scalar, otherwise return it."""
-
-  if isinstance(value, Array):
+def get_first(obj: TArrayTree) -> TArrayTree:
+  """Index the PyTree leaves `x` of `obj` by `x[0]` if they are not scalars."""
 
-    if value.ndim > 0:
-      return value[index]
-    else:
-      return value
+  def _get_first(value: Numeric) -> Numeric:
+    if isinstance(value, Array):
 
-  elif isinstance(value, (float, int)):
-    return value
-  raise ValueError("The input should be an instance of `Numeric`.")
+      if value.ndim > 0:
+        if jax.config.jax_pmap_shmap_merge:
+          shard_data = value.addressable_shards[0].data
+          if not value.sharding.is_fully_replicated:
+            return shard_data.squeeze(0)
+          return shard_data
+        return value[0]
+      else:
+        return value
 
+    elif isinstance(value, (float, int)):
+      return value
+    raise ValueError("The input should be an instance of `Numeric`.")
 
-def get_first(obj: TArrayTree) -> TArrayTree:
-  """Index the PyTree leaves `x` of `obj` by `x[0]` if they are not scalars."""
-  return jax.tree_util.tree_map(index_if_not_scalar, obj)
+  return jax.tree_util.tree_map(_get_first, obj)
 
 
 def get_mean(obj: TArrayTree) -> TArrayTree: