Internal change.

PiperOrigin-RevId: 872326871

Author: Orbax Authors

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

@@ -31,6 +31,15 @@ class methods:
     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
@@ -478,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:
@@ -491,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

checkpoint/orbax/checkpoint/_src/arrays/fragments_test.py

@@ -927,6 +927,72 @@ def test_abstract_fragments_of(self, value_fn):
         fragment_t(index=np.s_[0:2:1, 2:3:1]),
     ])
 
+  @parameterized.named_parameters(
+      ('np_array', NpFragments),
+      ('jnp_array', JaxFragments),
+  )
+  def test_abstract_fragments_like_concrete_fragments(
+      self, template_fragments_t: FragmentsT
+  ):
+    fragments_t = AbstractFragments
+    shape = (2, 3)
+    dtype = np.dtype(np.float32)
+    fragment_t = fragments_t.FRAGMENT_T
+
+    template_np_api = template_fragments_t.FRAGMENT_T.NP_API
+    template = template_fragments_t(
+        shape=shape,
+        dtype=dtype,
+        fragments=[
+            template_fragments_t.FRAGMENT_T(
+                index=np.s_[0:2:1, 0:1:1], value=template_np_api.ones((2, 1))
+            ),
+            template_fragments_t.FRAGMENT_T(
+                index=np.s_[0:2:1, 2:3:1], value=template_np_api.ones((2, 1))
+            ),
+        ],
+    )
+
+    fs = fragments_t.like(template)
+
+    self.assertEqual(fs.shape, shape)
+    self.assertEqual(fs.dtype, dtype)
+    self.assertEqual(fs.fragments, [
+        fragment_t(index=np.s_[0:2:1, 0:1:1], value=None),
+        fragment_t(index=np.s_[0:2:1, 2:3:1], value=None),
+    ])
+
+  @parameterized.named_parameters(
+      ('np_array', NpFragments),
+      ('jnp_array', JaxFragments),
+  )
+  def test_concrete_fragments_like_abstract_fragments(
+      self, fragments_t: FragmentsT
+  ):
+    shape = (2, 3)
+    dtype = np.dtype(np.float32)
+    fragment_t = fragments_t.FRAGMENT_T
+    np_api = fragment_t.NP_API
+
+    template = AbstractFragments(
+        shape=shape,
+        dtype=dtype,
+        fragments=[
+            AbstractFragment(index=np.s_[0:2:1, 0:1:1]),
+            AbstractFragment(index=np.s_[0:2:1, 2:3:1]),
+        ],
+    )
+    value = np_api.ones(shape, dtype=dtype)
+
+    fs = fragments_t.like(template, value)
+
+    self.assertEqual(fs.shape, shape)
+    self.assertEqual(fs.dtype, dtype)
+    self.assertEqual(fs.fragments, [
+        fragment_t(index=np.s_[0:2:1, 0:1:1], value=value[0:2:1, 0:1:1]),
+        fragment_t(index=np.s_[0:2:1, 2:3:1], value=value[0:2:1, 2:3:1]),
+    ])
+
 
 @parameterized.named_parameters(
     ('abstract_fragments', AbstractFragments),