def METHOD_NAME(self, batch_size, **kwargs):
    # get batched shuffled indices
    # for example: batch_size=2; self.group=5
    # indices = [
    #     [0, 2, 3, 4, 1],
    #     [4, 1, 0, 2, 3]
    # ]
    indices_distribution = self._random_generator.random_uniform(
        (batch_size, self.groups)
    )
    indices = tf.argsort(indices_distribution, axis=-1)
    return indices