Internal change

PiperOrigin-RevId: 868740843

Author: Orbax Authors

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

@@ -615,22 +615,6 @@ def _gen_distinct_addressable_indices(
   yield from (np_utils.from_hashable_index(i) for i in distinct_indices)
 
 
-def abstract_fragments(
-    x: jax.Array | jax.ShapeDtypeStruct | FS,
-) -> AbstractFragments:
-  """Returns abstract fragments matching the given object."""
-  if isinstance(x, AbstractFragments):
-    return x
-  else:
-    if isinstance(x, _GenericFragments):
-      indices = (fragment.index for fragment in x.fragments)
-    else:
-      indices = addressable_shards(x)
-    return AbstractFragments(x.shape, x.dtype, [
-        AbstractFragment(index=index) for index in indices
-    ])
-
-
 def validate_fragments_can_be_stacked(fragments: FSconcrete) -> None:
   """Validates that the given fragments can be stacked."""
   if not fragments.fragments:

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

@@ -1118,59 +1118,6 @@ def test_unsharded_array_is_fully_replicated(self):
     )
 
 
-class AbstractFragmentsTest(parameterized.TestCase):
-
-  def test_returns_abstract_fragments_instance_itself(
-      self
-  ):
-    fragments = AbstractFragments(
-        shape=(2, 3), dtype=np.dtype(np.float32), fragments=[]
-    )
-    self.assertIs(fragments, array_fragments.abstract_fragments(fragments))
-
-  @parameterized.named_parameters(
-      ('np_fragments', NpFragments),
-      ('jax_fragments', JaxFragments),
-  )
-  def test_converts_concrete_fragments(
-      self, fragments_t: ConcreteFragmentsT
-  ):
-    fragment_t = fragments_t.FRAGMENT_T
-    np_api = fragment_t.NP_API
-    concrete_fragments = fragments_t(
-        shape=(2, 3),
-        dtype=np.dtype(np.float32),
-        fragments=[
-            fragment_t(index=np.s_[0:2:1, 0:3:1], value=np_api.arange(6)),
-        ],
-    )
-    expected_abstract_fragments = AbstractFragments(
-        shape=(2, 3),
-        dtype=np.dtype(np.float32),
-        fragments=[
-            AbstractFragment(
-                index=np.s_[0:2:1, 0:3:1], value=None
-            ),
-        ],
-    )
-    self.assertEqual(
-        expected_abstract_fragments,
-        array_fragments.abstract_fragments(concrete_fragments),
-    )
-
-  def test_converts_fully_replicated_shape_dtype_struct(self):
-    self.assertEqual(
-        AbstractFragments(
-            shape=(4, 5),
-            dtype=np.dtype(np.float32),
-            fragments=[AbstractFragment(index=np.s_[0:4:1, 0:5:1])],
-        ),
-        array_fragments.abstract_fragments(
-            jax.ShapeDtypeStruct((4, 5), np.dtype(np.float32))
-        ),
-    )
-
-
 class StackFragmentsTest(parameterized.TestCase):
 
   @parameterized.named_parameters(