code: def tree_flatten_with_names(
      tree: PyTree,
  ) -> Tuple[List[Tuple[str, jnp.ndarray]], PyTree]:
      """Populates tree_flatten with leaf names.
  
      This function populates output of tree_flatten with leaf names, using a
      custom traversal that produces names is provided. The custom traversal does
      NOT have to traverse tree in the same order as jax, as we take care of
      automatically aligning jax' and custom traversals.
  
      Args:
        tree: python tree.
  
      Returns:
        A list of values with names: [(name, value), ...]
      """
      vals, tree_def = jax.tree_util.tree_flatten(tree)
  
      # "Fake" token tree that is used to track Jax internal tree traversal and
      # adjust our custom tree traversal to be compatible with it.
      tokens = range(len(vals))
      token_tree = tree_def.unflatten(tokens)
      val_names, perm = zip(*_traverse_with_names(token_tree))
      inv_perm = np.argsort(perm)
  
      # Custom traversal should visit the same number of leaves.
      assert len(val_names) == len(vals)
  
      return [
          (val_names[i], v) for i, v in zip(inv_perm, vals)
      ], tree_def
  
imports: import copy
  import dataclasses
  import re
  from typing import Any, Callable, Generator, List, Optional, Tuple, Union
  from absl import logging
  import flax
  import jax
  import jax.numpy
  import ml_collections
  import numpy
  import optax
label_desc: Use image node_id 7 for calling a global function with example usage: tree_flatten_with_names(tree) and returns: unknown,  tree_def