code: def get_shallow_tree(is_leaf, tree):
      """Returns a shallow tree, expanding only when is_leaf(subtree) is False."""
      return tree_util.tree_map(is_leaf, tree, is_leaf=is_leaf)
  
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 5 for calling a global function with example usage: get_shallow_tree(is_leaf, tree) and returns: tree_util