code: def stack_forest(forest):
      """Helper function to stack the leaves of a sequence of pytrees.
  
      Args:
        forest: a sequence of pytrees (e.g tuple or list) of matching structure
          whose leaves are arrays with individually matching shapes.
      Returns:
        A single pytree of the same structure whose leaves are individually
          stacked arrays.
      """
      stack_args = lambda *args: np.stack(args)
      return jax.tree_util.tree_map(stack_args, *forest)
  
file_path: flax/flax/training/common_utils.py
imports: import jax
  import jax.numpy
  import numpy
  from jax import lax
label_desc: Use image node_id 3 for calling a global function with example usage: stack_forest(forest) and returns: jax