trnrand 0.4.0

Random number generation for AWS Trainium via NKI

trnrand

Random number generation for AWS Trainium via NKI.

Seeded pseudo-random distributions, quasi-random sequences for quasi-Monte Carlo, and on-device Philox RNG targeting the GpSimd engine.

Part of the trnsci scientific computing suite (github.com/trnsci).

Current phase

trnrand follows the trnsci 5-phase roadmap. Active work is tracked in phase-labeled GitHub issues:

• Phase 1 — correctness (active): Philox 4×32 and Box-Muller NKI kernels scaffolded with CPU-reference oracles; awaiting hardware validation on trn1 / trn2.
• Phase 3 — perf: batched-tile RNG streaming, NEFF cache reuse.
• Phase 4 — multi-chip: stream-partitioned RNG across NeuronCores.
• Phase 5 — generation: trn2 PSUM sizing for larger batched tile generation.

(No Phase 2 for trnrand — the precision story is inherited where relevant.)

Suite-wide tracker: trnsci/trnsci#1.

Install

pip install trnrand

# With Neuron hardware support
pip install trnrand[neuron]

Usage

import trnrand

# Seeded, reproducible generation
g = trnrand.Generator(seed=42)

# Standard distributions
x = trnrand.normal(1000, mean=0.0, std=1.0, generator=g)
u = trnrand.uniform(1000, low=-1.0, high=1.0, generator=g)
e = trnrand.exponential(1000, rate=2.0, generator=g)

# Quasi-random sequences (better convergence for MC integration)
sobol_pts = trnrand.sobol(1024, n_dims=5, seed=42)
halton_pts = trnrand.halton(1024, n_dims=3)
lhs_pts = trnrand.latin_hypercube(100, n_dims=4)

# Module-level seeding
trnrand.manual_seed(42)
x = trnrand.standard_normal(256)

Operations

Category	Function	Description
Distributions	uniform	U[low, high)
	normal	N(μ, σ²)
	standard_normal	N(0, 1)
	exponential	Exp(λ)
	bernoulli	Bernoulli(p)
	randint	Uniform integers [low, high)
	randperm	Random permutation
	truncated_normal	Bounded normal (rejection sampling)
Quasi-random	sobol	Sobol sequence (scrambled)
	halton	Halton sequence
	latin_hypercube	Latin Hypercube Sampling

MC vs QMC Example

python examples/mc_integration.py

Compares pseudo-random vs Sobol quasi-random for estimating the volume of a 5-D hypersphere. QMC converges O(1/N) vs O(1/√N).

Status

• Seeded Generator with state management
• Standard distributions (uniform, normal, exponential, Bernoulli, etc.)
• Sobol, Halton, Latin Hypercube sequences
• MC vs QMC integration example
• NKI Philox kernel on GpSimd
• On-device Box-Muller (uniform → normal)
• Benchmarks vs cuRAND

Related Projects

Project	What
trnfft	FFT + complex ops
trnblas	BLAS operations
trnsolver	Linear solvers

License

Apache 2.0 — Copyright 2026 Scott Friedman

Disclaimer

trnsci is an independent open-source project. It is not sponsored by, endorsed by, or affiliated with Amazon.com, Inc., Amazon Web Services, Inc., or Annapurna Labs Ltd.

"AWS", "Amazon", "Trainium", "Inferentia", "NeuronCore", "Neuron SDK", and related identifiers are trademarks of their respective owners and are used here solely for descriptive and interoperability purposes. Use does not imply endorsement, partnership, or any other relationship.

All work, opinions, analyses, benchmark results, architectural commentary, and editorial judgments in this repository and on trnsci.dev are those of the project's contributors. They do not represent the views, positions, or commitments of Amazon, AWS, or Annapurna Labs.

Feedback directed at the Neuron SDK or Trainium hardware is good-faith ecosystem commentary from independent users. It is not privileged information, is not pre-reviewed by AWS, and should not be read as authoritative about product roadmap, behavior, or quality.

For official AWS guidance, see aws-neuron documentation and the AWS Trainium product page.