Random

Deterministic RNG utilities for TorchLean (seed-threaded, pure).

Why not IO.rand / runtime randomness?

Lean (and mathlib) can generate random numbers via IO, but that gives effectful randomness whose results depend on hidden runtime state. For TorchLean, that is a poor fit:

• it breaks the "one semantics" contract for lowering and verification (graphs stop being a pure mathematical object unless you model the RNG state explicitly),
• it makes replays and certificate checking depend on hidden runtime state,
• it complicates typed graph recording (effects and mixed dtypes).

Instead we use a deterministic pseudorandom generator and treat randomness as a deterministic function of an explicit seed (and a counter/stream id). This mirrors the JAX/functional RNG style and keeps the semantic core pure.

What this file provides

• a small 64-bit PRNG (SplitMix64-style mixing),
• deterministic sampling utilities keyed by (seed, counter, linearIndex),
• a deterministic way to build dropout-style {0,1} masks as shape-indexed tensors.

The module lives in the spec layer because the IR reference semantics (NN.IR.Semantics) needs randUniform and bernoulliMask nodes to denote pure tensors. Runtime code calls these same definitions directly; Session-level stochastic layers (for example TorchLean.Session.dropout) store RNG state in NatRefs and use these generators to build masks reproducibly.

If you want PyTorch-like randomness at the boundary, prefer sampling an initial seed in IO and then using the seeded RNG from that point onward (TorchLean.Session.initRngFromIO).

SplitMix64-style mixing

def Spec.Random.splitmix64 (x : UInt64) : UInt64

SplitMix64-style mixing function on 64-bit words.

This is used as a compact deterministic PRNG core: we treat "randomness" as a pure function of an explicit seed/counter/index.

def Spec.Random.keyOf (seed counter : ℕ) : UInt64

Derive a per-call key from a (seed, counter) pair.

def Spec.Random.nextSeed (seed counter : ℕ) : ℕ

Advance the seed deterministically once per RNG use.

Sampling helpers

def Spec.Random.sampleNat (key : UInt64) (linearIndex : ℕ) (denom : ℕ := 2 ^ 32) : ℕ

Deterministic "random" natural number in [0, denom) when denom > 0, derived from key and a linear index.

def Spec.Random.sampleUnit {α : Type} [Context α] (u denom : ℕ) : α

Convert u/denom into α using backend arithmetic.

Even when u < denom, floating-point rounding can make the result equal to 1.

def Spec.Random.keepBit {α : Type} [Context α] (keepProb : α) (u denom : ℕ) : α

Decide whether to keep an element given keepProb and a sample u ∈ [0, denom).

Handle the probability endpoints before converting the sample: binary32 can round the largest 32-bit draw to 1, but a keep probability of 1 must still keep every element.

Uniform tensors

def Spec.Random.Internal.uniform {α : Type} [TorchLean.Storage α] [Context α] (key : UInt64) {s : Shape} : ℕ → TorchLean.Tensor α s

Fill a tensor with deterministic unit draws.

Each coordinate gets its own counter value. The integer draw depends only on key and the linear coordinate; conversion and division use the selected backend. Reusing those inputs on the same backend reproduces the result, which can round to the upper endpoint 1.

def Spec.Random.uniform {α : Type} [TorchLean.Storage α] [Context α] (key : UInt64) {s : Shape} : TorchLean.Tensor α s

Build a uniform tensor over the whole shape, starting the deterministic stream at offset 0.

Mask construction

def Spec.Random.Internal.mask {α : Type} [TorchLean.Storage α] [Context α] (key : UInt64) (keepProb : α) {s : Shape} : ℕ → TorchLean.Tensor α s

Fill a tensor with a Bernoulli keep mask, one for kept coordinates and zero for dropped.

def Spec.Random.mask {α : Type} [TorchLean.Storage α] [Context α] (key : UInt64) (keepProb : α) {s : Shape} : TorchLean.Tensor α s

Build a dropout-style mask over the whole shape, starting the stream at offset 0.

Standard normal tensors

def Spec.Random.boxMullerCos {α : Type} [Context α] (u1 u2 : α) : α

Box-Muller transform: turn two independent uniforms u1,u2 ∈ (0,1) into a standard normal sample.

We return only the cos branch:

z = sqrt(-2 * log u1) * cos(2π * u2).

Notes:

• We clamp u1 below by ε to avoid log 0.
• This is intended as a deterministic pseudo-normal sampler for examples and benchmarking. It is not a cryptographic RNG.

Reference:

• Box & Muller (1958), "A Note on the Generation of Random Normal Deviates".

def Spec.Random.normalScalar {α : Type} [Context α] (key : UInt64) (linearIndex : ℕ) : α

Deterministic standard normal N(0,1) sample derived from key and a linear index.

We use two 32-bit uniforms (via sampleNat) per output scalar and apply the Box-Muller transform.

def Spec.Random.Internal.normal {α : Type} [TorchLean.Storage α] [Context α] (key : UInt64) {s : Shape} : ℕ → TorchLean.Tensor α s

Fill a tensor with independent standard normal draws, coordinate by coordinate.

def Spec.Random.normal {α : Type} [TorchLean.Storage α] [Context α] (key : UInt64) {s : Shape} : TorchLean.Tensor α s

Build a standard-normal tensor over the whole shape, starting the stream at offset 0.