code: def tree_fn(fn, **unmapped_kwargs):
      """Wrap a function to jax.tree_map over its arguments.
  
      You may set some named arguments via a partial to skip the `tree_map` on those
      arguments. Usual caveats of `partial` apply (e.g. set named args must be a
      suffix of the argument list).
  
      Args:
        fn: the function to be wrapped.
        **unmapped_kwargs: the named arguments to be set via a partial.
  
      Returns:
        a function
      """
      pfn = functools.partial(fn, **unmapped_kwargs)
  
      def _wrapped(*args):
          return jax.tree_map(pfn, *args)
  
      return _wrapped
  
file_path: rlax/rlax/_src/tree_util.py
imports: import functools
  from typing import Any, Callable, Sequence
  import chex
  import jax
  import jax.numpy
  import numpy
  from rlax._src import base
label_desc: Use image node_id 6 for calling a global function with example usage: tree_fn(fn) and returns: _wrapped, jax