code: def all_gather(tree: Any):
      """Gather across different hosts and flatten the first two dimensions."""
      gather_flat = lambda x: jnp.concatenate(
          jax.lax.all_gather(x, "batch"), 0
      )
      return jax.tree_util.tree_map(gather_flat, tree)
  
file_path: scenic/scenic/projects/ncr/utils.py
imports: import os
  from typing import Any, Dict, Optional, Tuple
  from flax import jax_utils
  from flax.training import checkpoints
  import jax
  import jax.numpy
  import numpy
  from scenic.train_lib import train_utils
  from tensorflow.io import gfile
label_desc: Use image node_id 4 for calling a global function with example usage: all_gather(tree) and returns: jax