Randomness¶

tinax.randomness owns typed scalar key validation, ordered deterministic coordinate derivation, and explicit continuation-key ownership. It never discovers process identity or stores mutable RNG state.

Concrete Python coordinates must be between 0 and 2**32 - 1. Dynamic coordinates inside JIT-compiled calls must be scalar unsigned JAX values of at most 32 bits, preventing silent fold_in collisions from signed or oversized values.

See the Deterministic Randomness guide for task-oriented usage.

import jax

from tinax.randomness import derive_process_step_key, split_key

key = derive_process_step_key(jax.random.key(0), process_index=0, step=10)
next_key, batch_keys = split_key(key, count=8)

derive_key ¶

derive_key(key: Array, *coordinates: int | Array) -> Array

Fold ordered integer coordinates into one typed scalar key without consuming global state.

Parameters:

Name	Type	Description	Default
`key`	`Array`	Typed scalar JAX PRNG key to derive from.	required
`*coordinates`	`int \| Array`	Integer coordinates folded in order via `jax.random.fold_in`. Concrete Python values must lie in `0..2**32 - 1`; JIT-time values must be unsigned scalar JAX arrays no wider than 32 bits.	`()`

Returns:

Type	Description
`Array`	A new typed scalar JAX key derived from `key` and the coordinates. The
`Array`	input key is not consumed.

Raises:

Type	Description
`TypeError`	If `key` is not a typed scalar PRNG key, or a coordinate is a boolean, non-scalar, or a JAX scalar with a signed or oversized dtype.
`ValueError`	If `key` is not scalar, or a Python coordinate falls outside `0..2**32 - 1`.

derive_process_step_key ¶

derive_process_step_key(
    key: Array,
    *,
    process_index: int | Array,
    step: int | Array,
    stream: int | Array = 0,
) -> Array

Derive one key using the stable process, step, then stream coordinate order.

Parameters:

Name	Type	Description	Default
`key`	`Array`	Typed scalar JAX PRNG key to derive from.	required
`process_index`	`int \| Array`	Coordinate identifying the process; folded in first.	required
`step`	`int \| Array`	Coordinate identifying the step; folded in second.	required
`stream`	`int \| Array`	Coordinate identifying the substream; folded in last. Defaults to 0.	`0`

Returns:

Type	Description
`Array`	A new typed scalar JAX key derived in process, step, then stream order.

Raises:

Type	Description
`TypeError`	If `key` is not a typed scalar PRNG key, or a coordinate is a boolean, non-scalar, or a JAX scalar with a signed or oversized dtype.
`ValueError`	If `key` is not scalar, or a Python coordinate falls outside `0..2**32 - 1`.

split_key ¶

split_key(key: Array, *, count: int) -> tuple[Array, Array]

Return a continuation key and a leading-axis array of explicitly consumed operation keys.

Parameters:

Name	Type	Description	Default
`key`	`Array`	Typed scalar JAX PRNG key to split.	required
`count`	`int`	Number of operation keys to produce. Must be a non-negative integer.	required

Returns:

Type	Description
`Array`	A `(continuation_key, operation_keys)` tuple. `continuation_key` is a
`Array`	scalar key for further derivation; `operation_keys` has leading dimension
`tuple[Array, Array]`	`count`.

Raises:

Type	Description
`TypeError`	If `key` is not a typed scalar PRNG key, or `count` is not an integer (booleans are rejected).
`ValueError`	If `key` is not scalar, or `count` is negative.