Function split

jax/_src/random/core.py:319–331

Splits a PRNG key into `num` new keys by adding a leading axis. Args: key: a PRNG key (from ``key``, ``split``, ``fold_in``). num: optional, a positive integer (or tuple of integers) indicating the number (or shape) of keys to produce. Defaults to 2. Returns: An array-like ob

(key: ArrayLike, num: int | tuple[int, ...] = 2)

  return prng.random_split(key, shape=shape)

def split(key: ArrayLike, num: int | tuple[int, ...] = 2) -> Array:
  """Splits a PRNG key into `num` new keys by adding a leading axis.

  Args:
    key: a PRNG key (from ``key``, ``split``, ``fold_in``).
    num: optional, a positive integer (or tuple of integers) indicating
      the number (or shape) of keys to produce. Defaults to 2.

  Returns:
    An array-like object of `num` new PRNG keys.
  """
  typed_key, wrapped = _check_prng_key("split", key)
  return _return_prng_keys(wrapped, _split(typed_key, num))


def _key_impl(keys: Array) -> PRNGImpl: