[pmap] Handle edge cases in get_from_first_device for jax_pmap_shmap_merge.

pmap replication/sharding semantics are not consistent, so different users of the acme library do different things. The real solution is to migrate away from pmap entirely, but IIUC, there may not be staffing for this.

Add a robust _unreplicate helper function that handles:
- 0-dimensional scalar arrays (cannot be indexed)
- SingleDeviceSharding or single local device (use simple indexing)
- Fully replicated arrays (return addressable shard data without squeeze)
- Default sharded case (squeeze the leading replicated dimension)

This prevents errors when processing nested structures containing arrays
with different sharding configurations under jax_pmap_shmap_merge mode.

PiperOrigin-RevId: 864435972
Change-Id: I6e151cbe0ed771f62762ec4f2d8e11a972e2a4ac

Author: danielsuo

acme/jax/utils.py

@@ -387,9 +387,21 @@ def get_from_first_device(nest: N, as_numpy: bool = True) -> N:
   # 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:
-    zeroth_nest = jax.tree_util.tree_map(
-        lambda x: x.addressable_shards[0].data.squeeze(0), nest
-    )
+
+    def _unreplicate(x):
+      if hasattr(x, 'ndim') and x.ndim == 0:
+        return x
+      if not hasattr(x, 'sharding') or isinstance(
+          x.sharding, jax.sharding.SingleDeviceSharding
+      ):
+        return x
+      if 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)
+
+    zeroth_nest = jax.tree_util.tree_map(_unreplicate, nest)
   else:
     zeroth_nest = jax.tree_util.tree_map(lambda x: x[0], nest)
   return jax.device_get(zeroth_nest) if as_numpy else zeroth_nest