include/tenzr/modules/random/random.h¶

// 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.

#ifndef TENZR_MODULES_RANDOM_RANDOM_H_
#define TENZR_MODULES_RANDOM_RANDOM_H_

#include <cmath>
#include <cstdint>
#include <vector>

#include "tenzr/common/export.h"
#include "tenzr/frontend/dtype.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 {

  // Stateless JAX-style Functional PRNG Distributions

  [[nodiscard]] TENZR_API Array Uniform(const PRNGKey& key, const Shape& shape, float low = 0.0f, float high = 1.0f, DType dtype = DType::kFloat32);

  [[nodiscard]] TENZR_API Array Normal(const PRNGKey& key, const Shape& shape, float mean = 0.0f, float stddev = 1.0f, DType dtype = DType::kFloat32);

  [[nodiscard]] TENZR_API Array Bernoulli(const PRNGKey& key, const Shape& shape, float p = 0.5f);

  [[nodiscard]] TENZR_API Array Rademacher(const PRNGKey& key, const Shape& shape, DType dtype = DType::kFloat32);

  [[nodiscard]] TENZR_API Array Categorical(const PRNGKey& key, const Array& logits, const Shape& shape);

  // Neural Network Weight Initializers (Functional PRNGKey overloads)

  [[nodiscard]] TENZR_API Array GlorotUniform(const PRNGKey& key, const Shape& shape, int64_t fan_in, int64_t fan_out, DType dtype = DType::kFloat32);

  [[nodiscard]] TENZR_API Array GlorotNormal(const PRNGKey& key, const Shape& shape, int64_t fan_in, int64_t fan_out, DType dtype = DType::kFloat32);

  [[nodiscard]] TENZR_API Array KaimingUniform(const PRNGKey& key, const Shape& shape, int64_t fan_in, DType dtype = DType::kFloat32);

  [[nodiscard]] TENZR_API Array KaimingNormal(const PRNGKey& key, const Shape& shape, int64_t fan_in, DType dtype = DType::kFloat32);

  // Legacy Stateful Generator Utilities (Backwards Compatibility)

  void Seed(uint64_t s);
  void Seed();

  Tensor Uniform(Shape shape, float low = 0.0f, float high = 1.0f, DType dtype = DType::kFloat32, Generator& gen = DefaultGenerator());
  Tensor Normal(Shape shape, float mean = 0.0f, float stddev = 1.0f, DType dtype = DType::kFloat32, Generator& gen = DefaultGenerator());
  Tensor Bernoulli(Shape shape, float p = 0.5f, Generator& gen = DefaultGenerator());

  Tensor GlorotUniform(Shape shape, int64_t fan_in, int64_t fan_out, DType dtype, Generator& gen = DefaultGenerator());
  Tensor GlorotUniform(Shape shape, int64_t fan_in, int64_t fan_out, Generator& gen = DefaultGenerator());
  Tensor GlorotNormal(Shape shape, int64_t fan_in, int64_t fan_out, DType dtype, Generator& gen = DefaultGenerator());
  Tensor GlorotNormal(Shape shape, int64_t fan_in, int64_t fan_out, Generator& gen = DefaultGenerator());
  Tensor KaimingUniform(Shape shape, int64_t fan_in, DType dtype, Generator& gen = DefaultGenerator());
  Tensor KaimingUniform(Shape shape, int64_t fan_in, Generator& gen = DefaultGenerator());
  Tensor KaimingNormal(Shape shape, int64_t fan_in, DType dtype, Generator& gen = DefaultGenerator());
  Tensor KaimingNormal(Shape shape, int64_t fan_in, Generator& gen = DefaultGenerator());

  inline Tensor Rand(Shape shape, Generator& gen = DefaultGenerator()) { return Uniform(std::move(shape), 0.0f, 1.0f, DType::kFloat32, gen); }

  inline Tensor Rand(Shape shape, float low, float high, Generator& gen = DefaultGenerator()) {
    return Uniform(std::move(shape), low, high, DType::kFloat32, gen);
  }

  inline Tensor Randn(Shape shape, Generator& gen = DefaultGenerator()) { return Normal(std::move(shape), 0.0f, 1.0f, DType::kFloat32, gen); }

  inline Tensor Randn(Shape shape, float mean, float stddev, Generator& gen = DefaultGenerator()) {
    return Normal(std::move(shape), mean, stddev, DType::kFloat32, gen);
  }

}  // namespace tenzr::random

#endif  // TENZR_MODULES_RANDOM_RANDOM_H_