code: def map_tree(fn, tree, *args):
      """Maps `fn` over the leaves of a nested structure."""
      return tree_util.tree_map(fn, tree, *args)
  
imports: import functools
  import jax
  from jax import lax
  from jax import random
  from jax import tree_util
  import jax.numpy
label_desc: Use image node_id 1 for calling a global function with example usage: map_tree(fn, tree) and returns: tree_util