code: def tree_leaves(tree, ignore_nones=True):
      """Gets the leaves of a tree."""
  
      # Right now this is just `tree_flatten`, but we keep this separate since
      # JAX's tree_flatten returns the structure of the tree as well.
      flattened = tree_flatten(tree)
      return [
          flat
          for flat in flattened
          if (not ignore_nones) or flat is not None
      ]
  
file_path: trax/trax/fastmath/numpy.py
imports: import numpy
  from scipy.special import logsumexp
  from trax.shapes import signature
label_desc: Use image node_id 12 for calling a global function with example usage: tree_leaves(tree, ignore_nones) and returns: unknown