code: def _jax_batched_split(key, n):
      """A version of jax.random.split` that works PRNG keys of any rank."""
      split = lambda k: jax.random.split(k, n)
      for _ in range(key.ndim):
          split = jax.vmap(split, out_axes=1)
      return split(key)
  
file_path: probability/tensorflow_probability/python/internal/loop_util.py
imports: import numpy
  import tensorflow.compat.v1
  import tensorflow.compat.v2
  from tensorflow_probability.python.internal import prefer_static
  from tensorflow_probability.python.internal import tensorshape_util
  from tensorflow.python.ops import control_flow_util
label_desc: Use image node_id 1 for calling a global function with example usage: _jax_batched_split(key, n) and returns: split