code: def _traverse_with_names(
      tree: PyTree,
  ) -> Generator[Tuple[str, PyTree], None, None]:
      """Traverses nested dicts/dataclasses and emits (leaf_name, leaf_val)."""
      if dataclasses.is_dataclass(tree):
          tree = flax.serialization.to_state_dict(tree)
      if isinstance(tree, (dict, flax.core.frozen_dict.FrozenDict)):
          keys = sorted(tree.keys())
          for key in keys:
              for path, v in _traverse_with_names(tree[key]):
                  yield (key + "/" + path).rstrip("/"), v
      else:
          yield "", tree
  
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 6 for calling a global function with example usage: _traverse_with_names(tree) without return types