code: def get_shallow_tree(is_leaf, tree):
      """Returns a shallow tree, expanding only when is_leaf(subtree) is False."""
      return nest.get_traverse_shallow_structure(
          lambda t: not is_leaf(t), tree
      )
  
imports: import functools
  import numpy
  import six
  import tensorflow.compat.v2
  from tensorflow.python.util import nest
label_desc: Use image node_id 5 for calling a global function with example usage: get_shallow_tree(is_leaf, tree) and returns: nest