Skip to content

← Flax NNX step 2 of 100

Medium Primitives

nnx.Rngs Container

Jump to the task

Why this matters

JAX random functions take an explicit key. Using the same key with the same sampling operation repeats the same pseudorandom sample. To obtain a new draw, you need a different key.

nnx.Rngs manages key generation with stateful streams. Each stream stores a base key and a counter. Calling the stream derives a key by folding the counter into the base key, then advances the counter. This is not the same key sequence as repeatedly calling jax.random.split.

The container handles this bookkeeping, but sampling is still performed by JAX. NNX code can also use explicit keys when that suits the task.

Worked example

import jax
import jax.numpy as jnp
from flax import nnx

rngs = nnx.Rngs(42)
key1 = rngs.params()
key2 = rngs.params()
a = jax.random.normal(key1, (4,), dtype=jnp.float32)
b = jax.random.normal(key2, (4,), dtype=jnp.float32)
result = a + b

assert result.shape == (4,)
assert not jnp.array_equal(jax.random.key_data(key1), jax.random.key_data(key2))

reset = nnx.Rngs(42)
assert jnp.array_equal(
    jax.random.key_data(reset.params()), jax.random.key_data(key1)
)

Reconstructing the container resets the sequence. Reusing the existing container advances it. Reproducibility depends on the seed, stream call order, random-key implementation and sampling configuration. A seed alone is not a guarantee of bitwise agreement across every backend or version.

Default and named streams

nnx.Rngs(0) creates a default stream. With no explicit params or dropout stream, both rngs.params() and rngs.dropout() use that same default stream, advancing the same counter. Changing their call order changes which key each operation receives.

nnx.Rngs(params=0, dropout=1) creates two named streams with separate counters. A draw from dropout does not advance params. If a requested name is absent and there is no default stream, the lookup fails.

Different names alone do not make different sequences: named streams initialized from the same seed can produce matching keys at matching counter positions. Choose distinct seeds or suitably derived keys when the streams should generate different draws.

The container is not itself a PRNG key. Obtain a key by calling a stream; rngs() can also draw from the default stream when one exists.

Problem

Implement rngs_two_keys_sum(seed, k1_offset, k2_offset):

  1. Construct rngs = nnx.Rngs(int(seed)) once.
  2. Obtain key1 and key2 with two consecutive rngs.params() calls.
  3. Draw a float32 standard-normal vector of shape (4,) from each key.
  4. Add k1_offset to the first vector and k2_offset to the second.
  5. Return the elementwise sum of the two shifted vectors.

Seed is a nonnegative integer-valued scalar that may arrive as a float. The offsets are finite scalar values. Return a float32 vector of shape (4,). Keep offsets as JAX-compatible values; Python float conversions are unnecessary and break tracing when the offsets are dynamic.

Use the specified stream calls to match the fixtures. Splitting a single key or drawing one vector of shape (2, 4) uses a different random operation and need not produce the same numbers.

Common pitfalls

  • Do not reconstruct Rngs before the second draw. That resets the stream and repeats its first key.
  • Do not reuse key1 for the second sample. That repeats the first vector instead of producing the intended second draw.
  • Do not mistake fallback names for separate streams. An explicit default stream is shared by all missing names that fall back to it.
  • Separate pseudorandom draws do not mean every sampled number must differ. Check the prescribed key sequence, not an assumed absence of collisions.
  • If applying JIT to this setup function, mark seed static because the constructor uses its Python integer value. The offsets can stay dynamic.