Internal

PiperOrigin-RevId: 842451645

Author: Grain Team

grain/_src/core/constants.py

@@ -25,6 +25,10 @@
 # Index of the dataset from which to take the record. Only present when mixing.
 # Starts with 0.
 DATASET_INDEX = "_dataset_index"
+# Index of the dataset group from which to take the dataset. A group is defined
+# by a set of datasets having identical transformations. Only present when
+# mixing. Starts with 0.
+DATASET_GROUP_INDEX = "_dataset_group_index"
 # Epoch for the record. When mixing datasets this is the epoch over the dataset,
 # not the mixture. Starts with 1.
 EPOCH = "_epoch"
@@ -35,5 +39,5 @@
 RECORD = "_record"
 
 META_FEATURES = frozenset(
-    [INDEX, RECORD_KEY, DATASET_INDEX, EPOCH, SEED, RECORD]
+    [INDEX, RECORD_KEY, DATASET_INDEX, DATASET_GROUP_INDEX, EPOCH, SEED, RECORD]
 )