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: an eager session counts the random tensors it draws (rngCounter) and derives each one's key from that count, so its masks are reproducible.

If you want PyTorch-like randomness at the boundary, sample an initial seed in IO once and pass it as the runtime seed; everything after that point is seeded.

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 : ℕ) : α

Round the exact sampling fraction into the selected scalar type.

Forming the fraction before casting avoids overflowing either integer in a low-precision format. Rounding can still send a draw below 1 to the upper endpoint 1. A zero denominator retains the backend's division behavior; generators use a positive denominator.

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} (linearOffset : ℕ) : 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; its exact fraction is cast to 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} (linearOffset : ℕ) : 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 the backend's default safeguard. Avoiding log 0 requires that this safeguard be positive and that the backend comparison select it; Context alone gives no such law.
• 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 approximate standard-normal 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. The finite sampling grid, clamping, and backend arithmetic do not give an exact N(0,1) law.

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

Fill a tensor with approximate normal draws keyed by their linear coordinates.

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

Build an approximate normal tensor over the whole shape, starting the stream at offset 0.