// Copyright 2026 The Tenzr Project Authors
//
// SPDX-License-Identifier: Apache-2.0
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// https://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
#include "tenzr/modules/random/random.h"
#include <algorithm>
#include <cmath>
#include <cstdint>
#include <cstring>
#include <random>
#include <utility>
#include <vector>
#include "tenzr/apiutils/apiutils.h"
#include "tenzr/common/error.h"
#include "tenzr/frontend/dtype.h"
#include "tenzr/frontend/memory_pool.h"
#include "tenzr/frontend/shape.h"
#include "tenzr/frontend/tensor.h"
#include "tenzr/modules/random/generator.h"
#include "tenzr/modules/random/prng_key.h"
namespace tenzr::random {
// PRNGKey Implementation
PRNGKey::PRNGKey(const Array& arr) {
TENZR_CHECK(arr.shape().NumElements() >= 2, "PRNGKey array must contain at least 2 elements");
auto vec = arr.Astype(DType::kUInt32).ToVector<uint32_t>();
k0 = vec[0];
k1 = vec[1];
}
Array PRNGKey::ToArray() const {
Array a = Array::Zeros(Shape({2}), DType::kUInt32);
a.SetData(std::vector<uint32_t>{k0, k1});
return a;
}
PRNGKey Key(uint64_t seed_val) { return PRNGKey(seed_val); }
// Threefry 2x32 Counter-based PRNG Engine (20 Rounds)
namespace {
[[nodiscard]] inline constexpr uint32_t rotl32(uint32_t v, uint32_t d) noexcept { return (v << d) | (v >> (32 - d)); }
constexpr uint32_t kThreefryC240 = 0x1BD11BDAU;
constexpr uint32_t kRotations0[4] = {13, 15, 26, 6};
constexpr uint32_t kRotations1[4] = {17, 29, 16, 24};
} // namespace
std::pair<uint32_t, uint32_t> Threefry2x32(uint32_t k0, uint32_t k1, uint32_t c0, uint32_t c1) noexcept {
uint32_t ks[3];
ks[0] = k0;
ks[1] = k1;
ks[2] = k0 ^ k1 ^ kThreefryC240;
uint32_t x0 = c0 + ks[0];
uint32_t x1 = c1 + ks[1];
// 5 blocks of 4 rounds = 20 rounds
// Block 0: rounds 0-3
for (uint32_t r : kRotations0) {
x0 += x1;
x1 = rotl32(x1, r) ^ x0;
}
x0 += ks[1];
x1 += ks[2] + 1U;
// Block 1: rounds 4-7
for (uint32_t r : kRotations1) {
x0 += x1;
x1 = rotl32(x1, r) ^ x0;
}
x0 += ks[2];
x1 += ks[0] + 2U;
// Block 2: rounds 8-11
for (uint32_t r : kRotations0) {
x0 += x1;
x1 = rotl32(x1, r) ^ x0;
}
x0 += ks[0];
x1 += ks[1] + 3U;
// Block 3: rounds 12-15
for (uint32_t r : kRotations1) {
x0 += x1;
x1 = rotl32(x1, r) ^ x0;
}
x0 += ks[1];
x1 += ks[2] + 4U;
// Block 4: rounds 16-19
for (uint32_t r : kRotations0) {
x0 += x1;
x1 = rotl32(x1, r) ^ x0;
}
x0 += ks[2];
x1 += ks[0] + 5U;
return {x0, x1};
}
SplitKeys Split(const PRNGKey& k) {
auto keys = Split(k, 2);
return {keys[0], keys[1]};
}
std::vector<PRNGKey> Split(const PRNGKey& k, size_t num) {
TENZR_CHECK(num > 0, "split num must be greater than 0");
std::vector<PRNGKey> result;
result.reserve(num);
for (size_t i = 0; i < num; ++i) {
uint32_t c0 = static_cast<uint32_t>(2 * i);
uint32_t c1 = static_cast<uint32_t>(2 * i + 1);
auto [k0_new, k1_new] = Threefry2x32(k.k0, k.k1, c0, c1);
result.emplace_back(k0_new, k1_new);
}
return result;
}
PRNGKey FoldIn(const PRNGKey& k, uint32_t data) {
auto [k0_new, k1_new] = Threefry2x32(k.k0, k.k1, data, 0x5A5A5A5AU);
return PRNGKey(k0_new, k1_new);
}
std::vector<uint32_t> ThreefryRandomBits(const PRNGKey& k, size_t count) {
std::vector<uint32_t> bits;
bits.reserve(count);
size_t pairs = (count + 1) / 2;
for (size_t i = 0; i < pairs; ++i) {
uint32_t c0 = static_cast<uint32_t>(2 * i);
uint32_t c1 = static_cast<uint32_t>(2 * i + 1);
auto [b0, b1] = Threefry2x32(k.k0, k.k1, c0, c1);
bits.push_back(b0);
if (bits.size() < count) {
bits.push_back(b1);
}
}
return bits;
}
// Stateless JAX-style Functional PRNG Distributions
Array Uniform(const PRNGKey& key, const Shape& shape, float low, float high, DType dtype) {
TENZR_CHECK(low < high, "Uniform distribution requires low < high");
size_t count = static_cast<size_t>(shape.NumElements());
std::vector<uint32_t> bits = ThreefryRandomBits(key, count);
constexpr double kInv232 = 1.0 / 4294967296.0;
double range = static_cast<double>(high - low);
if (dtype == DType::kFloat32) {
std::vector<float> data(count);
for (size_t i = 0; i < count; ++i) {
double u = static_cast<double>(bits[i]) * kInv232;
data[i] = static_cast<float>(static_cast<double>(low) + u * range);
}
Array a = Array::Zeros(shape, DType::kFloat32);
a.SetData(data);
return a;
} else if (dtype == DType::kFloat64) {
std::vector<double> data(count);
for (size_t i = 0; i < count; ++i) {
double u = static_cast<double>(bits[i]) * kInv232;
data[i] = static_cast<double>(low) + u * range;
}
Array a = Array::Zeros(shape, DType::kFloat64);
a.SetData(data);
return a;
}
TENZR_THROW("Uniform distribution currently supports kFloat32 and Float64");
}
Array Normal(const PRNGKey& key, const Shape& shape, float mean, float stddev, DType dtype) {
TENZR_CHECK(stddev > 0.0f, "Normal distribution requires stddev > 0");
size_t count = static_cast<size_t>(shape.NumElements());
// For Box-Muller, generate even number of uniforms
size_t needed = (count % 2 == 0) ? count : (count + 1);
std::vector<uint32_t> bits = ThreefryRandomBits(key, needed);
constexpr double kTwoPi = 6.28318530717958647692;
constexpr double kInvDenom = 1.0 / 4294967298.0;
std::vector<float> data_f32;
std::vector<double> data_f64;
if (dtype == DType::kFloat32) {
data_f32.reserve(count);
} else {
data_f64.reserve(count);
}
for (size_t i = 0; i < needed; i += 2) {
double u1 = (static_cast<double>(bits[i]) + 1.0) * kInvDenom;
double u2 = (static_cast<double>(bits[i + 1]) + 1.0) * kInvDenom;
double r = std::sqrt(-2.0 * std::log(u1));
double theta = kTwoPi * u2;
double z0 = (r * std::cos(theta)) * static_cast<double>(stddev) + static_cast<double>(mean);
double z1 = (r * std::sin(theta)) * static_cast<double>(stddev) + static_cast<double>(mean);
if (dtype == DType::kFloat32) {
data_f32.push_back(static_cast<float>(z0));
if (data_f32.size() < count) {
data_f32.push_back(static_cast<float>(z1));
}
} else {
data_f64.push_back(z0);
if (data_f64.size() < count) {
data_f64.push_back(z1);
}
}
}
if (dtype == DType::kFloat32) {
Array a = Array::Zeros(shape, DType::kFloat32);
a.SetData(data_f32);
return a;
} else if (dtype == DType::kFloat64) {
Array a = Array::Zeros(shape, DType::kFloat64);
a.SetData(data_f64);
return a;
}
TENZR_THROW("Normal distribution currently supports kFloat32 and Float64");
}
Array Bernoulli(const PRNGKey& key, const Shape& shape, float p) {
TENZR_CHECK(p >= 0.0f && p <= 1.0f, "Bernoulli probability p must be in [0, 1]");
size_t count = static_cast<size_t>(shape.NumElements());
std::vector<uint32_t> bits = ThreefryRandomBits(key, count);
uint32_t threshold = static_cast<uint32_t>(static_cast<double>(p) * 4294967295.0);
std::vector<bool> data(count);
for (size_t i = 0; i < count; ++i) {
data[i] = (bits[i] <= threshold);
}
Array a = Array::Zeros(shape, DType::kBool);
a.SetData(data);
return a;
}
Array Rademacher(const PRNGKey& key, const Shape& shape, DType dtype) {
size_t count = static_cast<size_t>(shape.NumElements());
std::vector<uint32_t> bits = ThreefryRandomBits(key, count);
if (dtype == DType::kFloat32) {
std::vector<float> data(count);
for (size_t i = 0; i < count; ++i) {
data[i] = (bits[i] & 1U) ? 1.0f : -1.0f;
}
Array a = Array::Zeros(shape, DType::kFloat32);
a.SetData(data);
return a;
}
std::vector<double> data(count);
for (size_t i = 0; i < count; ++i) {
data[i] = (bits[i] & 1U) ? 1.0 : -1.0;
}
Array a = Array::Zeros(shape, DType::kFloat64);
a.SetData(data);
return a;
}
Array Categorical(const PRNGKey& key, const Array& logits, const Shape& shape) {
TENZR_CHECK(logits.shape().Rank() >= 1, "Categorical logits must have rank >= 1");
int64_t num_classes = logits.shape().Dim(logits.shape().Rank() - 1);
TENZR_CHECK(num_classes > 0, "Categorical num_classes must be positive");
size_t num_samples = static_cast<size_t>(shape.NumElements());
size_t total_gumbel = num_samples * static_cast<size_t>(num_classes);
std::vector<uint32_t> bits = ThreefryRandomBits(key, total_gumbel);
std::vector<float> logits_vec = logits.Astype(DType::kFloat32).ToVector<float>();
constexpr double kInvDenom = 1.0 / 4294967298.0;
std::vector<int32_t> out_classes(num_samples);
for (size_t s = 0; s < num_samples; ++s) {
int32_t best_c = 0;
float max_val = -1.0e30f;
for (int64_t c = 0; c < num_classes; ++c) {
size_t bit_idx = s * static_cast<size_t>(num_classes) + static_cast<size_t>(c);
double u = (static_cast<double>(bits[bit_idx]) + 1.0) * kInvDenom;
// Gumbel(0, 1) = -ln(-ln(u))
float g = static_cast<float>(-std::log(-std::log(u)));
float logit_val = logits_vec[static_cast<size_t>(c % static_cast<int64_t>(logits_vec.size()))];
float score = logit_val + g;
if (score > max_val) {
max_val = score;
best_c = static_cast<int32_t>(c);
}
}
out_classes[s] = best_c;
}
Array res = Array::Zeros(shape, DType::kInt32);
res.SetData(out_classes);
return res;
}
// Neural Network Weight Initializers (Functional PRNGKey overloads)
Array GlorotUniform(const PRNGKey& key, const Shape& shape, int64_t fan_in, int64_t fan_out, DType dtype) {
float limit = std::sqrt(6.0f / static_cast<float>(fan_in + fan_out));
return Uniform(key, shape, -limit, limit, dtype);
}
Array GlorotNormal(const PRNGKey& key, const Shape& shape, int64_t fan_in, int64_t fan_out, DType dtype) {
float stddev = std::sqrt(2.0f / static_cast<float>(fan_in + fan_out));
return Normal(key, shape, 0.0f, stddev, dtype);
}
Array KaimingUniform(const PRNGKey& key, const Shape& shape, int64_t fan_in, DType dtype) {
float limit = std::sqrt(3.0f / static_cast<float>(fan_in));
return Uniform(key, shape, -limit, limit, dtype);
}
Array KaimingNormal(const PRNGKey& key, const Shape& shape, int64_t fan_in, DType dtype) {
float stddev = std::sqrt(2.0f / static_cast<float>(fan_in));
return Normal(key, shape, 0.0f, stddev, dtype);
}
// Legacy Stateful Generator Utilities (Backwards Compatibility)
void Seed(uint64_t s) { DefaultGenerator().ManualSeed(s); }
void Seed() { DefaultGenerator().SeedRandom(); }
Tensor Uniform(Shape shape, float low, float high, DType dtype, Generator& gen) {
TENZR_CHECK(low < high, "Uniform distribution requires low < high");
size_t count = static_cast<size_t>(shape.NumElements());
size_t bytes = count * SizeOf(dtype);
auto buf = MemoryPool::Instance().Allocate(bytes);
if (dtype == DType::kFloat32) {
std::uniform_real_distribution<float> dist(low, high);
float* ptr = buf->As<float>();
for (size_t i = 0; i < count; ++i) {
ptr[i] = dist(gen.Engine());
}
} else if (dtype == DType::kFloat64) {
std::uniform_real_distribution<double> dist(static_cast<double>(low), static_cast<double>(high));
double* ptr = buf->As<double>();
for (size_t i = 0; i < count; ++i) {
ptr[i] = dist(gen.Engine());
}
} else {
TENZR_THROW("Uniform distribution currently supports kFloat32 and Float64");
}
auto t = Tensor::Parameter(std::move(shape), dtype);
t.BindData(std::move(buf));
return t;
}
Tensor Normal(Shape shape, float mean, float stddev, DType dtype, Generator& gen) {
TENZR_CHECK(stddev > 0.0f, "Normal distribution requires stddev > 0");
size_t count = static_cast<size_t>(shape.NumElements());
size_t bytes = count * SizeOf(dtype);
auto buf = MemoryPool::Instance().Allocate(bytes);
if (dtype == DType::kFloat32) {
std::normal_distribution<float> dist(mean, stddev);
float* ptr = buf->As<float>();
for (size_t i = 0; i < count; ++i) {
ptr[i] = dist(gen.Engine());
}
} else if (dtype == DType::kFloat64) {
std::normal_distribution<double> dist(static_cast<double>(mean), static_cast<double>(stddev));
double* ptr = buf->As<double>();
for (size_t i = 0; i < count; ++i) {
ptr[i] = dist(gen.Engine());
}
} else {
TENZR_THROW("Normal distribution currently supports kFloat32 and Float64");
}
auto t = Tensor::Parameter(std::move(shape), dtype);
t.BindData(std::move(buf));
return t;
}
Tensor Bernoulli(Shape shape, float p, Generator& gen) {
TENZR_CHECK(p >= 0.0f && p <= 1.0f, "Bernoulli probability p must be in [0, 1]");
size_t count = static_cast<size_t>(shape.NumElements());
size_t bytes = count * sizeof(float);
auto buf = MemoryPool::Instance().Allocate(bytes);
std::bernoulli_distribution dist(static_cast<double>(p));
float* ptr = reinterpret_cast<float*>(buf->Data());
for (size_t i = 0; i < count; ++i) {
ptr[i] = dist(gen.Engine()) ? 1.0f : 0.0f;
}
auto t = Tensor::Parameter(std::move(shape), DType::kFloat32);
t.BindData(std::move(buf));
return t;
}
Tensor GlorotUniform(Shape shape, int64_t fan_in, int64_t fan_out, DType dtype, Generator& gen) {
float limit = std::sqrt(6.0f / static_cast<float>(fan_in + fan_out));
return Uniform(std::move(shape), -limit, limit, dtype, gen);
}
Tensor GlorotUniform(Shape shape, int64_t fan_in, int64_t fan_out, Generator& gen) {
return GlorotUniform(std::move(shape), fan_in, fan_out, DType::kFloat32, gen);
}
Tensor GlorotNormal(Shape shape, int64_t fan_in, int64_t fan_out, DType dtype, Generator& gen) {
float stddev = std::sqrt(2.0f / static_cast<float>(fan_in + fan_out));
return Normal(std::move(shape), 0.0f, stddev, dtype, gen);
}
Tensor GlorotNormal(Shape shape, int64_t fan_in, int64_t fan_out, Generator& gen) {
return GlorotNormal(std::move(shape), fan_in, fan_out, DType::kFloat32, gen);
}
Tensor KaimingUniform(Shape shape, int64_t fan_in, DType dtype, Generator& gen) {
float limit = std::sqrt(3.0f / static_cast<float>(fan_in));
return Uniform(std::move(shape), -limit, limit, dtype, gen);
}
Tensor KaimingUniform(Shape shape, int64_t fan_in, Generator& gen) { return KaimingUniform(std::move(shape), fan_in, DType::kFloat32, gen); }
Tensor KaimingNormal(Shape shape, int64_t fan_in, DType dtype, Generator& gen) {
float stddev = std::sqrt(2.0f / static_cast<float>(fan_in));
return Normal(std::move(shape), 0.0f, stddev, dtype, gen);
}
Tensor KaimingNormal(Shape shape, int64_t fan_in, Generator& gen) { return KaimingNormal(std::move(shape), fan_in, DType::kFloat32, gen); }
} // namespace tenzr::random