Fix nbytes calculation for ReplicaSlices to consider slices.

PiperOrigin-RevId: 794485404

Author: Orbax Authors

checkpoint/orbax/checkpoint/_src/serialization/replica_slices.py

@@ -115,8 +115,22 @@ def __post_init__(self):
 
   @property
   def nbytes(self) -> int:
-    slice_nbytes = math.prod(self.local_shape) * self.dtype.itemsize
-    return slice_nbytes * len(self.replica_slices)
+    """Returns the total number of bytes for all replica slices."""
+    total_bytes = 0
+    for rslice in self.replica_slices:
+      if rslice.slice_args is None:
+        # No slicing, use the full local_shape
+        slice_nbytes = math.prod(self.local_shape) * self.dtype.itemsize
+      else:
+        # Replica-parallel, calculate bytes for the specific slice
+        slice_shape = list(self.local_shape)
+        axis = rslice.slice_args.axis
+        slice_shape[axis] = (
+            rslice.slice_args.limit_index - rslice.slice_args.start_index
+        )
+        slice_nbytes = math.prod(slice_shape) * self.dtype.itemsize
+      total_bytes += slice_nbytes
+    return total_bytes
 
   def to_fragments(self) -> fragments.Fragments:
     """Converts replica slices to fragments."""

checkpoint/orbax/checkpoint/_src/serialization/replica_slices_test.py

@@ -334,6 +334,109 @@ def test_transfer(self, partitioned, use_replica_parallel):
       # With single-replica we transfer a single slice for each shard.
       self.assertLen(rslices.replica_slices, num_partitions)
 
+  def test_nbytes_no_slicing(self):
+    rslice = replica_slices.ReplicaSlice(
+        index=(slice(None),),
+        unsliced_data=np.zeros((2, 3), dtype=np.float32),
+        slice_args=None,
+    )
+    rslices = replica_slices.ReplicaSlices(
+        global_shape=(2, 3),
+        local_shape=(2, 3),
+        sharding=jax.sharding.NamedSharding(
+            jax.sharding.Mesh(np.array(jax.devices()), ('x',)),
+            jax.sharding.PartitionSpec(),
+        ),
+        dtype=np.dtype(np.float32),
+        is_on_host=True,
+        replica_slices=[rslice],
+    )
+    expected_nbytes = 2 * 3 * np.dtype(np.float32).itemsize
+    self.assertEqual(rslices.nbytes, expected_nbytes)
+
+  def test_nbytes_with_slicing(self):
+    rslice = replica_slices.ReplicaSlice(
+        index=(slice(None),),
+        unsliced_data=jax.numpy.zeros((4, 5), dtype=np.int16),
+        slice_args=replica_slices.SliceArgs(
+            start_index=1, limit_index=3, axis=0
+        ),
+    )
+    rslices = replica_slices.ReplicaSlices(
+        global_shape=(4, 5),
+        local_shape=(4, 5),
+        sharding=jax.sharding.NamedSharding(
+            jax.sharding.Mesh(np.array(jax.devices()), ('x',)),
+            jax.sharding.PartitionSpec(),
+        ),
+        dtype=np.dtype(np.int16),
+        is_on_host=False,  # Set to False to allow slice_args
+        replica_slices=[rslice],
+    )
+    # Shape of the slice is (2, 5)
+    expected_nbytes = 2 * 5 * np.dtype(np.int16).itemsize
+    self.assertEqual(rslices.nbytes, expected_nbytes)
+
+  def test_nbytes_multiple_slices(self):
+    rslice1 = replica_slices.ReplicaSlice(
+        index=(slice(None),),
+        unsliced_data=jax.numpy.zeros((4, 5), dtype=np.int16),
+        slice_args=replica_slices.SliceArgs(
+            start_index=0, limit_index=2, axis=0
+        ),
+    )
+    rslice2 = replica_slices.ReplicaSlice(
+        index=(slice(None),),
+        unsliced_data=jax.numpy.zeros((4, 5), dtype=np.int16),
+        slice_args=replica_slices.SliceArgs(
+            start_index=2, limit_index=4, axis=0
+        ),
+    )
+    rslices = replica_slices.ReplicaSlices(
+        global_shape=(4, 5),
+        local_shape=(4, 5),
+        sharding=jax.sharding.NamedSharding(
+            jax.sharding.Mesh(np.array(jax.devices()), ('x',)),
+            jax.sharding.PartitionSpec(),
+        ),
+        dtype=np.dtype(np.int16),
+        is_on_host=False,  # Set to False to allow slice_args
+        replica_slices=[rslice1, rslice2],
+    )
+    # Shape of each slice is (2, 5)
+    expected_nbytes = (2 * 5 * np.dtype(np.int16).itemsize) * 2
+    self.assertEqual(rslices.nbytes, expected_nbytes)
+
+  def test_nbytes_mixed_slicing(self):
+    rslice1 = replica_slices.ReplicaSlice(
+        index=(slice(None),),
+        unsliced_data=jax.numpy.zeros((4, 5), dtype=np.int16),  # jax.Array
+        slice_args=None,
+    )
+    rslice2 = replica_slices.ReplicaSlice(
+        index=(slice(None),),
+        unsliced_data=jax.numpy.zeros((4, 5), dtype=np.int16),  # jax.Array
+        slice_args=replica_slices.SliceArgs(
+            start_index=0, limit_index=2, axis=1
+        ),
+    )
+    rslices = replica_slices.ReplicaSlices(
+        global_shape=(4, 5),
+        local_shape=(4, 5),
+        sharding=jax.sharding.NamedSharding(
+            jax.sharding.Mesh(np.array(jax.devices()), ('x',)),
+            jax.sharding.PartitionSpec(),
+        ),
+        dtype=np.dtype(np.int16),
+        is_on_host=False,  # Set to False to allow slice_args
+        replica_slices=[rslice1, rslice2],
+    )
+    # bytes for rslice1: (4 * 5)
+    # bytes for rslice2: (4 * 2)
+    expected_nbytes = (4 * 5 * np.dtype(np.int16).itemsize) + (
+        4 * 2 * np.dtype(np.int16).itemsize
+    )
+    self.assertEqual(rslices.nbytes, expected_nbytes)
 
 if __name__ == '__main__':
   absltest.main()