code: def map_tree_up_to(shallow, fn, tree, *rest):
      """`map_tree` with recursion depth defined by depth of `shallow`."""
  
      def wrapper(_, *rest):
          return fn(*rest)
  
      return tree_util.tree_map(wrapper, shallow, tree, *rest)
  
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 4 for calling a global function with example usage: map_tree_up_to(shallow, fn, tree) and returns: tree_util, fn