code: def tree_split_key(rng_key: Array, tree_like: Any):
      """Generate random keys for each leaf in a tree.
  
      Args:
        rng_key: a JAX pseudo random number generator key.
        tree_like: a nested structure.
  
      Returns:
        a new key, and a tree of keys with same shape as `tree_like`.
      """
      leaves, treedef = jax.tree_util.tree_flatten(tree_like)
      rng_key, *keys = jax.random.split(rng_key, num=len(leaves) + 1)
      return rng_key, jax.tree_util.tree_unflatten(treedef, keys)
  
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 3 for calling a global function with example usage: tree_split_key(rng_key, tree_like) and returns: rng_key,  jax