Fixed bug in DataLoader where sharding remainder was always dropped, ignoring ShardOptions.drop_remainder=False.

PiperOrigin-RevId: 934593946

Author: nithintatikonda2002-source

CHANGELOG.md

@@ -12,6 +12,7 @@ changes. Best viewed [here](https://google-grain.readthedocs.io/en/latest/change
 * Deprecations:
 
 * Bug fixes:
+  * Fixed bug in DataLoader where sharding remainder was dropped even when ShardOptions.drop_remainder=False.
 
 ## Grain 0.2.18 (June 17, 2026)
 

grain/_src/python/data_loader.py

@@ -125,7 +125,15 @@ def __init__(
     super().__init__(dataset.MapDataset.source(data_source))
     self._sampler = sampler
     self._shard_options = shard_options
-    self.length = self._sampler_size() // self._shard_options.shard_count
+    sampler_size = self._sampler_size()
+    shard_count = self._shard_options.shard_count
+    if self._shard_options.drop_remainder:
+      self.length = sampler_size // shard_count
+    else:
+      remainder = sampler_size % shard_count
+      self.length = sampler_size // shard_count + (
+          1 if self._shard_options.shard_index < remainder else 0
+      )
 
   def _sampler_size(self) -> int:
     """Returns the length of the sampler."""

grain/_src/python/data_loader_test.py

@@ -825,6 +825,54 @@ def test_start_prefetch(
       data_loader_iterator.start_prefetch()
     self.assertEqual(list(data_loader_iterator), list(range(16)))
 
+  def test_data_loader_sharding_drop_remainder(self):
+    range_data_source = RangeDataSource(start=0, stop=10, step=1)
+
+    # Shard 0, drop_remainder=False
+    sampler = samplers.SequentialSampler(
+        num_records=len(range_data_source),
+        shard_options=sharding.ShardOptions(
+            shard_index=0, shard_count=3, drop_remainder=False
+        ),
+    )
+    dataloader = data_loader_lib.DataLoader(
+        data_source=range_data_source,
+        sampler=sampler,
+    )
+    actual = list(dataloader)
+    self.assertEqual(actual, [0, 3, 6, 9])
+    self.assertLen(actual, 4)
+
+    # Shard 0, drop_remainder=True
+    sampler = samplers.SequentialSampler(
+        num_records=len(range_data_source),
+        shard_options=sharding.ShardOptions(
+            shard_index=0, shard_count=3, drop_remainder=True
+        ),
+    )
+    dataloader = data_loader_lib.DataLoader(
+        data_source=range_data_source,
+        sampler=sampler,
+    )
+    actual = list(dataloader)
+    self.assertEqual(actual, [0, 3, 6])
+    self.assertLen(actual, 3)
+
+    # Shard 1, drop_remainder=False
+    sampler = samplers.SequentialSampler(
+        num_records=len(range_data_source),
+        shard_options=sharding.ShardOptions(
+            shard_index=1, shard_count=3, drop_remainder=False
+        ),
+    )
+    dataloader = data_loader_lib.DataLoader(
+        data_source=range_data_source,
+        sampler=sampler,
+    )
+    actual = list(dataloader)
+    self.assertEqual(actual, [1, 4, 7])
+    self.assertLen(actual, 3)
+
 
 class PyGrainDatasetIteratorTest(absltest.TestCase):