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] )