code: def all_gather(
      x,
      *,
      mesh: jax.sharding.Mesh,
      axis_name: str | Sequence[str],
      memory_space: pltpu.TPUMemorySpace = pltpu.VMEM,
  ):
      if isinstance(axis_name, str):
          axis_name = (axis_name,)
      # TODO(sharadmv): enable all gather over multiple axes
      if len(axis_name) > 1:
          raise NotImplementedError("Only one axis supported.")
      (axis_name,) = axis_name
      if mesh.shape[axis_name] == 1:
          # We can short-circuit here if our axis size is 1
          return x
  
      def ag_local(x_shard):
          axis_size = lax.psum(1, axis_name)
          out_shape = jax.ShapeDtypeStruct(
              (axis_size, *x_shard.shape), x_shard.dtype
          )
          out = pl.pallas_call(
              functools.partial(
                  ag_kernel, axis_name=axis_name, mesh=mesh
              ),
              out_shape=out_shape,
              mosaic_params=dict(collective_id=0),
              grid_spec=pltpu.PrefetchScalarGridSpec(
                  num_scalar_prefetch=0,
                  scratch_shapes=(
                      (
                          pltpu.SemaphoreType.DMA,
                          pltpu.SemaphoreType.DMA,
                      ),
                      (
                          pltpu.SemaphoreType.DMA,
                          pltpu.SemaphoreType.DMA,
                      ),
                  ),
                  in_specs=[pl.BlockSpec(memory_space=memory_space)],
                  out_specs=pl.BlockSpec(memory_space=memory_space),
              ),
          )(x_shard)
          return out.reshape(
              (axis_size * x_shard.shape[0], *x_shard.shape[1:])
          )
  
      return shard_map.shard_map(
          ag_local,
          mesh=mesh,
          in_specs=P(axis_name),
          out_specs=P(None),
          check_rep=False,
      )(x)
  
file_path: jax/jax/experimental/pallas/ops/tpu/all_gather.py
imports: from __future__ import annotations
  import functools
  from collections.abc import Sequence
  import jax
  from jax import lax
  from jax.experimental import pallas
  from jax.experimental import shard_map
  from jax.experimental.pallas import tpu
  import jax.numpy
label_desc: Use image node_id 3 for calling a global function with example usage: all_gather(x) and returns: shard_map, x, out