code: def tree_map_zipped(fn: Callable[..., Any], nests: Sequence[Any]):
      """Map a function over a list of identical nested structures.
  
      Args:
        fn: the function to map; must have arity equal to `len(list_of_nests)`.
        nests: a list of identical nested structures.
  
      Returns:
        a nested structure whose leaves are outputs of applying `fn`.
      """
      if not nests:
          return nests
      tree_def = tree_structure(nests[0])
      if any(tree_structure(x) != tree_def for x in nests[1:]):
          raise ValueError(
              "All elements must share the same tree structure."
          )
      return jax.tree_util.tree_unflatten(
          tree_def,
          [fn(*d) for d in zip(*[jax.tree_leaves(x) for x in nests])],
      )
  
file_path: rlax/rlax/_src/tree_util.py
imports: import functools
  from typing import Any, Callable, Sequence
  import chex
  import jax
  import jax.numpy
  import numpy
  from rlax._src import base
label_desc: Use image node_id 2 for calling a global function with example usage: tree_map_zipped(fn, nests) and returns: jax, nests