From 2644c6b31b528b28feb9c4cd8d63c13cedb6c784 Mon Sep 17 00:00:00 2001 From: Grain Team Date: Tue, 9 Dec 2025 16:54:04 -0800 Subject: [PATCH] Internal PiperOrigin-RevId: 842451645 --- grain/_src/core/constants.py | 6 +++++- 1 file changed, 5 insertions(+), 1 deletion(-) diff --git a/grain/_src/core/constants.py b/grain/_src/core/constants.py index ef3d3f57d..50cd7b660 100644 --- a/grain/_src/core/constants.py +++ b/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] )