def apply_on_global_batch(fn, **kwargs):
  """Transforms a fn to operate on global batch with (num_dev, dev_bs, ...)."""

  def new_fn(*input_seq):
    input_set = []
    for an_input in input_seq:
      num_dev, dev_bs = jax.tree_util.tree_leaves(an_input)[0].shape[:2]
      inputs = jax.tree_util.tree_map(
          lambda x: x.reshape(-1, *x.shape[2:]), an_input)
      input_set.append(inputs)
    outputs = fn(*input_set, **kwargs)
    outputs = jax.tree_util.tree_map(
        lambda x: x.reshape(num_dev, dev_bs, *x.shape[1:]), outputs)
    return outputs

  return new_fn
Transforms a fn to operate on global batch with (num_dev, dev_bs, ...).
videogvt/train_lib/train_utils.py