code: def tree_map_rngs(fn, tree):
      """Needed for mapping JAX random.* functions over PRNGKey leaves."""
      return jax.tree_util.tree_map(
          fn,
          tree,
          is_leaf=lambda x: hasattr(x, "dtype")
          and jax.dtypes.issubdtype(x.dtype, jax.dtypes.prng_key),
      )
  
imports: import collections
  import dataclasses
  import functools
  from typing import Any, Callable, Dict, Generic, Iterable, List, Mapping, Optional, Sequence, Tuple, TypeVar, Union
  import warnings
  from flax import traceback_util
  import jax
  from jax import random
  from .None import axes_scan, meta
  from .frozen_dict import freeze, unfreeze
  from .scope import CollectionFilter, DenyList, Filter, PRNGSequenceFilter, Scope, group_collections, in_filter, intersect_filters, is_filter_empty, subtract_filters, union_filters
label_desc: Use image node_id 1 for calling a global function with example usage: tree_map_rngs(fn, tree) and returns: jax