Internal change.

PiperOrigin-RevId: 712462394
google · Jan 13, 2025 · 86a8f1b · 86a8f1b
1 parent 10fdd0d
commit 86a8f1b
Show file tree

Hide file tree

Showing 4 changed files with 755 additions and 624 deletions.
diff --git a/grain/_src/python/dataset/transformations/BUILD b/grain/_src/python/dataset/transformations/BUILD
@@ -57,6 +57,17 @@ py_test(
     deps = ["//grain/_src/python/dataset"],
 )
 
+py_library(
+    name = "testing_util",
+    testonly = 1,
+    srcs = ["testing_util.py"],
+    srcs_version = "PY3",
+    deps = [
+        ":packing",
+        "//grain/_src/python/dataset",
+    ],
+)
+
 py_library(
     name = "flatmap",
     srcs = ["flatmap.py"],
@@ -116,6 +127,7 @@ py_test(
     srcs_version = "PY3",
     deps = [
         ":packing",
+        ":testing_util",
         "//grain/_src/python/dataset",
     ],
 )

diff --git a/grain/_src/python/dataset/transformations/packing.py b/grain/_src/python/dataset/transformations/packing.py
@@ -16,6 +16,7 @@
 from collections.abc import Sequence
 import copy
 from typing import Any, Optional
+from absl import logging
 from grain._src.python.dataset import dataset
 from grain._src.python.dataset import stats as dataset_stats
 from grain._src.python.dataset.transformations import packing_packed_batch
@@ -324,7 +325,6 @@ def __iter__(self) -> dataset.DatasetIterator:
         num_packing_bins=self._num_packing_bins,
         length_struct=self._length_struct,
         shuffle_bins=self._shuffle_bins,
-        shuffle_bins_group_by_feature=self._shuffle_bins_group_by_feature,
         meta_features=self._meta_features,
     )