Explicit random keys and reproducible draws

The second call can return exactly the same numbers

Call key.uniform(4) twice using the same key. Both tensors contain the same four values. Nothing advanced behind your back. In jax4j, PrngKey is an immutable seed-and-counter pair; drawing from it does not replace or mutate it. Reproducibility is explicit enough to become a visible data dependency.

That behavior is useful when replaying a simulation or comparing two algorithms with matched noise. It also means you must deliberately choose new keys or counter ranges when you want fresh draws. Merely repeating a function call is not a stream-management strategy.

Choose Run all, then change seeds, identifiers and shapes. This notebook uses the bundled jax4j 0.2.1 CPU backend. The examples print generated values and deterministic relationships; they do not claim measured statistical quality from these tiny samples. Reloading discards your Java state and edits.

1. Reusing a key repeats a request

A uniform draw returns F32 values in [0, 1). The request includes the key, distribution, shape and backend. Within this notebook's fixed implementation, repeating the same request reproduces the values exactly. A seed alone is not a complete record of an experiment if any of those other choices change.

Expect same values: true and root counter 0. Change the seed to 124 and rerun setup. The generated sequence changes, but its two repeated requests still agree. Return to 123 if you want the same printed values as an earlier run. Recording a seed makes a particular run repeatable; it does not eliminate sampling uncertainty in an estimator.

Try this: request shape [2, 2] instead of [4]. On this CPU uniform implementation, the same four counter positions fill the flattened output; the shape changes their organization. Print the shape alongside the values so reproducibility does not obscure an unintended layout change.

2. Reserve the next uniform counter range

For the CPU uniform generator in jax4j 0.2.1, a tensor with four entries uses four consecutive counter positions. advance(4) returns a new key beginning four positions later. The original key remains unchanged. We can compare this second block with the tail of a single eight-entry request.

Expect matching tail: true, then root=0, next=4. Advance by one instead and compare the two four-entry blocks: their counter ranges overlap in three positions. A different starting counter does not automatically imply a disjoint block of the size you plan to request.

Counter accounting depends on the distribution implementation. This version's CPU normal generator uses two counter positions per returned normal value. It applies a Box–Muller transformation to uniform values. A four-entry normal request therefore uses eight counter positions; advancing by four would not reserve a disjoint next block. Prefer derived keys when you do not want to couple stream management to such implementation details.

3. Derive named substreams by splitting

split(2) derives two child keys without consuming or mutating the parent. Repeating the same split produces the same child keys. Assigning a child to a specific purpose makes that purpose's randomness visible: for example, one for initialization and another for perturbing observations.

Expect repeat split: true. The two child keys request different deterministic streams. Do not interpret four printed values as proof of statistical independence or distribution quality. The useful guarantee shown here is reproducible derivation: the same parent and derivation procedure recover the same children.

Splitting is not a hidden advancing iterator. Calling root.split(2) again for the next batch simply recreates those children. Keep a deliberate hierarchy of purposes, steps or identifiers instead of repeatedly asking the root for “new” children with the same arguments.

4. Tie randomness to a stable identifier

foldIn(id) deterministically mixes an identifier into a key. This lets an observation receive the same random request even if you process it in another order. The next cell visits identifiers 7, 2, then 7 again. The two lines for identifier seven should contain exactly the same values.

Try reordering the identifiers to [2, 7, 2]. The values associated with each identifier remain stable. Use a persistent observation identifier if you want that property; a position in a shuffled list names the position, not the observation. For step-specific randomness, derive a key from both step and observation, using a consistent order such as root.foldIn(step).foldIn(id).

The order of derivations matters. Do not assume folding in step then identifier is interchangeable with identifier then step. Give different uses their own branches so an observation's augmentation does not accidentally reuse the same request as its diagnostic noise. Keep the key-derivation convention with the rest of the experiment configuration.

5. A captured key stays fixed during replay

uniformLike generates in an existing tensor's backend, device and transformation context. The key itself remains static Java data. Closing over the same key in a compiled function therefore records the same random request for each replay; it does not produce a fresh seed per call.

Expect same replay: true and one plan. Change the input values and the additive result changes, but the random component still comes from that fixed key and shape. When you need fresh noise for each application, generate it outside the compiled function using a deliberately derived key and pass the noise tensor as an explicit input leaf.

A PrngKey record is not automatically a dynamic compiled-tree argument. Its seed and counter are Java longs, while this version's compiled trees require tensor leaves. Changing a Java key in a closure also does not invalidate old plans. Use the capture rules from Chapter 3 when deciding where generation and key derivation belong.

6. Mapping does not split a fixed key

A fixed key used inside a mapped body requests one logical noise vector. That unmapped random vector is shared when it is added to each mapped row. To allocate different uniform counter positions to each row, generate a whole [2, 3] tensor outside the mapped function, then pass or use it explicitly.

The shared result repeats its first three values in the second row. The whole-batch draw uses counter positions zero through five: its first row matches that three-value request, and its second row uses the next three positions. Those statements describe counter allocation; they do not promise that every drawn value must differ from every other value.

Try shape [3, 3] in both inputs, updating the zero array to length nine. The shared request repeats across three rows, while the batch draw allocates nine positions. For stable per-observation draws across reordered batches, derive keys from identifiers as above rather than relying only on row position.

Return to mapped and shared arguments to feed a precomputed noise tensor alongside your data. Revisit capture rules when a replayed random computation is unexpectedly identical. The key, shape and mapping arrangement together explain exactly which request you made.