code: def tree_map_with_names(
      f: Callable[[jnp.ndarray], jnp.ndarray],
      param_tree: PyTree,
      match_name_fn: Callable[[str], bool] = lambda name: True,
  ) -> PyTree:
      """Like jax.tree_util.tree_map but with a filter on the leaf path name.
  
      Args:
        f: The function to be applied to each parameter in `param_tree`. Takes value
          as argument.
        param_tree: The tree of parameters `f` should be applied to.
        match_name_fn: This function is called with each tree leaf's path name,
          which has a path-like format ("a/b/c"), and decides whether `f` should be
          applied to that leaf or the leaf should be kept as-is.
  
      Returns:
        A tree identical in structure to `param_tree` but with the leaves the
        result of calling `f` on them in the cases where `match_name_fn` returns
        True for that leaf's path name.
      """
      names_and_vals, tree_def = tree_flatten_with_names(param_tree)
      vals = [
          f(v) if match_name_fn(name) else v
          for name, v in names_and_vals
      ]
      return tree_def.unflatten(vals)
  
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 8 for calling a global function with example usage: tree_map_with_names(f, param_tree, match_name_fn) and returns: tree_def