Internal change

PiperOrigin-RevId: 868682897

Author: Orbax Authors

checkpoint/orbax/checkpoint/_src/arrays/fragments.py

@@ -24,12 +24,28 @@ class methods:
     FS is `AbstractFragments` then x may additionally be a
     `jax.ShapeDtypeStruct`.
   - `none_of(x)`: gives an FS shaped like x, with no fragments.
+  - `addressable_shards_of(x)`: gives an FS shaped like x, with one fragment for
+    each distinct addressable shard of x. If FS is `AbstractFragments` then x
+    may additionally be a `jax.ShapeDtypeStruct`.
+  - `of(x, *indices)`: gives an FS shaped like x, with one fragment for each
+    index in `indices`. If FS is concrete then the fragment values will be
+    slices of x. If FS is `AbstractFragments` then x may additionally be
+    a `jax.ShapeDtypeStruct`.
+  - `like(fragments, x)` takes an existing Fragments instance and returns a new
+    instance of FS, with the same shape and dtype and fragment indices, but
+    with the fragment values replaced by slices of x. If FS is
+    `AbstractFragments` then x must be `None` and may be omitted.
+
+Use `FS.of()` and friends to make instances out of arraylike things.
+Use `FS.like()` to convert between different kinds of Fragments.
+Use `{np, jnp}.asarray(fragments)` to make arraylike things out of (full)
+Fragments.
 """
 # TODO(b/465196209): Remove when support for Python 3.10 is dropped.
 from __future__ import annotations
 
 import dataclasses
-from typing import Any, ClassVar, Generic, Literal, Sequence, TypeAlias, TypeVar
+from typing import Any, ClassVar, Generator, Generic, Literal, Sequence, TypeAlias, TypeVar
 
 import jax
 import numpy as np
@@ -369,6 +385,16 @@ def none_of(cls: type[FS], x: Any) -> FS:
     """Returns a Fragments with no fragments."""
     return cls._of(x, indices=[])
 
+  @classmethod
+  def addressable_shards_of(cls: type[FS], x: Any) -> FS:
+    """Returns a Fragments exactly spanning the distinct addressable shards of `x`."""
+    return cls._of(x, indices=_gen_distinct_addressable_indices(x))
+
+  @classmethod
+  def of(cls: type[FS], x: Any, *, indices: Sequence[Index]) -> FS:
+    """Returns a Fragments exactly spanning the given indices."""
+    return cls._of(x, indices=indices)
+
   def is_degenerate(self) -> bool:
     """Whether this contains only degenerate fragments."""
     return all(f.is_degenerate() for f in self.fragments)
@@ -461,11 +487,26 @@ def _of(cls: type[FS], x: Any, *, indices: Sequence[Index]) -> FS:
     fragments = [cls.FRAGMENT_T(index=index) for index in indices]
     return cls(x.shape, x.dtype, fragments)
 
+  @classmethod
+  def like(
+      cls: type[FS],
+      fragments: _GenericFragments[Any],
+      value: Literal[None] = None,
+  ) -> FS:
+    del value
+    return cls(
+        shape=fragments.shape,
+        dtype=fragments.dtype,
+        fragments=[
+            cls.FRAGMENT_T(index=f.index) for f in fragments.fragments
+        ],
+    )
+
 
 @dataclasses.dataclass(frozen=True, init=False)
-class NpFragments(_GenericFragments[NpFragment]):
-  """A collection of fragments whose values are of type `np.ndarray`."""
-  FRAGMENT_T = NpFragment
+class _ConcreteFragments(_GenericFragments[Fconcrete]):
+  """A collection of concrete fragments."""
+  FRAGMENT_T: ClassVar[type[Fconcrete]]  # The type of fragment values.
 
   @classmethod
   def _of(cls: type[FS], x: Any, *, indices: Sequence[Index]) -> FS:
@@ -474,19 +515,37 @@ def _of(cls: type[FS], x: Any, *, indices: Sequence[Index]) -> FS:
     fragments = [cls.FRAGMENT_T(index=i, value=x[i]) for i in indices]
     return cls(x.shape, x.dtype, fragments)
 
+  @classmethod
+  def like(
+      cls: type[FS], fragments: _GenericFragments[Any], value: Aconcrete
+  ) -> FS:
+    _check_fragment_value_type(value, cls.FRAGMENT_T.ARRAY_T)
+    if fragments.shape != value.shape or fragments.dtype != value.dtype:
+      raise ValueError(
+          f'Fragments type {fragments.dtype}[{fragments.shape}] does'
+          f' not match value type {value.dtype}[{value.shape}].'
+      )
+    return cls(
+        shape=fragments.shape,
+        dtype=fragments.dtype,
+        fragments=[
+            cls.FRAGMENT_T(index=f.index, value=value[f.index])
+            for f in fragments.fragments
+        ],
+    )
+
 
 @dataclasses.dataclass(frozen=True, init=False)
-class JaxFragments(_GenericFragments[JaxFragment]):
+class NpFragments(_ConcreteFragments[NpFragment]):
+  """A collection of fragments whose values are of type `np.ndarray`."""
+  FRAGMENT_T = NpFragment
+
+
+@dataclasses.dataclass(frozen=True, init=False)
+class JaxFragments(_ConcreteFragments[JaxFragment]):
   """A collection of fragments whose values are of type `jax.Array`."""
   FRAGMENT_T = JaxFragment
 
-  @classmethod
-  def _of(cls: type[FS], x: Any, *, indices: Sequence[Index]) -> FS:
-    """Returns a Fragments with one fragment for each index."""
-    _check_fragment_value_type(x, cls.FRAGMENT_T.ARRAY_T)
-    fragments = [cls.FRAGMENT_T(index=i, value=x[i]) for i in indices]
-    return cls(x.shape, x.dtype, fragments)
-
 
 # Extra names for backwards compatibility. Most loading and saving code still
 # wants to deal with NumPy arrays so that views and operations on them
@@ -534,6 +593,28 @@ def addressable_shards(x: jax.Array | jax.ShapeDtypeStruct) -> list[Index]:
   ]
 
 
+def _gen_distinct_addressable_indices(
+    x: np.ndarray | jax.Array | jax.ShapeDtypeStruct,
+) -> Generator[Index, None, None]:
+  """Yields fragment indices for distinct addressable shards of x."""
+  match x:
+    case jax.Array() | jax.ShapeDtypeStruct():
+      if not x.sharding:
+        raise ValueError(
+            'Cannot determine addressable shards of jax.ShapeDtypeStruct with'
+            ' no sharding.'
+        )
+      indices = addressable_shards(x)
+    case np.ndarray():
+      indices = (tuple(slice(0, dim, 1) for dim in x.shape),)
+    case _:
+      raise TypeError(f'Unsupported type: {type(x)}')
+  distinct_indices = sorted({
+      *(np_utils.to_hashable_index(i, shape=x.shape) for i in indices)
+  })
+  yield from (np_utils.from_hashable_index(i) for i in distinct_indices)
+
+
 def abstract_fragments(
     x: jax.Array | jax.ShapeDtypeStruct | FS,
 ) -> AbstractFragments: