Repository: google/jax

Issue:
Multihost gather operation
Currently there exists `jax.experimental.multihost_utils.process_allgather`, but no `process_gather`.

I am running the script on TPU v4-16.

Example:

```python
import jax
from jax import Array
import jax.numpy as jnp
from jax.sharding import Mesh, NamedSharding, PartitionSpec as P
import numpy as np
from types import EllipsisType

def shard_array_to_multihost(arr: Array, axis: int | EllipsisType) -> Array:
    shape = arr.shape
    devices: np.ndarray = np.array(jax.devices())

    if axis is ...:
        mesh = Mesh(devices, ('a',))
        sharding = NamedSharding(mesh, P(None))
    else:
        sharding_tuple_ = [1] * len(shape)
        sharding_tuple_[axis] = -1
        sharding_tuple = tuple(sharding_tuple_)
        name_tuple = tuple('abcdefghijklmnopqrstuvwxyz'[:len(shape)])

        mesh = Mesh(devices.reshape(sharding_tuple), name_tuple)
        sharding = NamedSharding(mesh, P(*name_tuple))

    xs = [jax.device_put(arr[i], device) for device, i in sharding.addressable_devices_indices_map(shape).items()]
    return jax.make_array_from_single_device_arrays(shape, sharding, xs)

arr = jnp.arange(2 * 16 * 8).reshape((2, 16, 8))
arr = shard_array_to_multihost(arr, axis=2)

arr = jax.experimental.multihost_utils.process_allgather(arr)
if jax.process_index() == 0:
    print(arr)
```

However, what I want to achieve is:

```python
if jax.process_index() == 0:
    print(process_gather(arr))
```


Discussion:
Besides, `jax.experimental.multihost_utils.process_allgather` returns a Numpy array instead of a JAX array. Why?
What does process_gather mean? If it means to gather all the local shards of the array then you can just write that on your own using `arr.addressable_shards`. 
Thank you! However, I think I cannot find the documentation of `arr.addressable_shards`. It is not included in https://jax.readthedocs.io/en/latest/_autosummary/jax.Array.html
I guess we need to fix that! Thanks for letting me know.

Resolving patch:
diff --git a/jax/_src/basearray.py b/jax/_src/basearray.py
--- a/jax/_src/basearray.py
+++ b/jax/_src/basearray.py
@@ -16,7 +16,11 @@
 
 import abc
 import numpy as np
-from typing import Union
+from typing import Any, Sequence, Union
+
+# TODO(jakevdp): fix import cycles and define these.
+Shard = Any
+Sharding = Any
 
 # Array is a type annotation for standard JAX arrays and tracers produced by
 # core functions in jax.lax and jax.numpy; it is not meant to include
@@ -46,11 +50,64 @@ def f(x: Array) -> Array:  # type annotations are valid for traced and non-trace
 
   __slots__ = ['__weakref__']
 
-  # at property must be defined because we overwrite its docstring in
-  # lax_numpy.py
   @property
-  def at(self):
-    raise NotImplementedError("property must be defined in subclasses")
+  @abc.abstractmethod
+  def dtype(self) -> np.dtype:
+    """The data type (:class:`numpy.dtype`) of the array."""
+
+  @property
+  @abc.abstractmethod
+  def ndim(self) -> int:
+    """The number of dimensions in the array."""
+
+  @property
+  @abc.abstractmethod
+  def size(self) -> int:
+    """The total number of elements in the array."""
+
+  @property
+  @abc.abstractmethod
+  def shape(self) -> tuple[int, ...]:
+    """The shape of the array."""
+
+  # Documentation for sharding-related methods and properties defined on ArrayImpl:
+  @abc.abstractmethod
+  def addressable_data(self, index: int) -> "Array":
+    """Return an array of the addressable data at a particular index."""
+
+  @property
+  @abc.abstractmethod
+  def addressable_shards(self) -> Sequence[Shard]:
+    """List of addressable shards."""
+
+  @property
+  @abc.abstractmethod
+  def global_shards(self) -> Sequence[Shard]:
+    """List of global shards."""
+
+  @property
+  @abc.abstractmethod
+  def is_fully_addressable(self) -> bool:
+    """Is this Array fully addressable?
+
+    A jax.Array is fully addressable if the current process can address all of
+    the devices named in the :class:`Sharding`. ``is_fully_addressable`` is
+    equivalent to "is_local" in multi-process JAX.
+
+    Note that fully replicated is not equal to fully addressable i.e.
+    a jax.Array which is fully replicated can span across multiple hosts and is
+    not fully addressable.
+    """
+
+  @property
+  @abc.abstractmethod
+  def is_fully_replicated(self) -> bool:
+    """Is this Array fully replicated?"""
+
+  @property
+  @abc.abstractmethod
+  def sharding(self) -> Sharding:
+    """The sharding for the array."""
 
 
 Array.__module__ = "jax"
diff --git a/jax/_src/core.py b/jax/_src/core.py
--- a/jax/_src/core.py
+++ b/jax/_src/core.py
@@ -626,11 +626,22 @@ def check_bool_conversion(arr: Array, warn_on_empty=False):
                       "ambiguous. Use a.any() or a.all()")
 
 
+def _aval_property(name):
+  return property(lambda self: getattr(self.aval, name))
+
 
 class Tracer(typing.Array):
   __array_priority__ = 1000
   __slots__ = ['_trace', '_line_info']
 
+  dtype = _aval_property('dtype')
+  ndim = _aval_property('ndim')
+  size = _aval_property('size')
+  shape = _aval_property('shape')
+
+  def __init__(self, trace: Trace):
+    self._trace = trace
+
   def _error_repr(self):
     if self.aval is None:
       return f"traced array with aval {self.aval}"
@@ -655,9 +666,6 @@ def tobytes(self, order="C"):
       f"The tobytes() method was called on {self._error_repr()}."
       f"{self._origin_msg()}")
 
-  def __init__(self, trace: Trace):
-    self._trace = trace
-
   def __iter__(self):
     return iter(self.aval._iter(self))
 
diff --git a/jax/_src/numpy/array_methods.py b/jax/_src/numpy/array_methods.py
--- a/jax/_src/numpy/array_methods.py
+++ b/jax/_src/numpy/array_methods.py
@@ -240,6 +240,7 @@ def _view(arr: Array, dtype: Optional[DTypeLike] = None, type: None = None) -> A
 
 
 def _notimplemented_flat(self):
+  """Not implemented: Use :meth:`~jax.Array.flatten` instead."""
   raise NotImplementedError("JAX Arrays do not implement the arr.flat property: "
                             "consider arr.flatten() instead.")
 
@@ -800,5 +801,3 @@ def register_jax_array_methods():
   _set_array_attributes(ArrayImpl)
 
   _set_array_abstract_methods(Array)
-
-  Array.at.__doc__ = _IndexUpdateHelper.__doc__


Test patch: