examples/research_projects/jax-projects/big_bird/bigbird_flax.py/DataCollator/__call__
class DataCollator:    def __call__(self, batch):
        batch = self.collate_fn(batch)
        batch = jax.tree_util.tree_map(shard, batch)
        return batch