src/python/random_bindings.cc¶

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

#include "bindings.h"

namespace tenzr::python {

  void BindRandom(nb::module_& m) {
    auto rand_m = m.def_submodule("random", "Stateless Threefry2x32 PRNG and functional distributions");

    nb::class_<random::PRNGKey>(rand_m, "PRNGKey")
        .def(nb::init<uint64_t>(), "seed"_a = 0)
        .def_prop_ro("k0", [](const random::PRNGKey& k) { return k.k0; })
        .def_prop_ro("k1", [](const random::PRNGKey& k) { return k.k1; })
        .def("__repr__",
             [](const random::PRNGKey& k) {
               std::ostringstream ss;
               ss << "PRNGKey(" << k.k0 << ", " << k.k1 << ")";
               return ss.str();
             })
        .def("__eq__", [](const random::PRNGKey& a, const random::PRNGKey& b) { return a.k0 == b.k0 && a.k1 == b.k1; });

    rand_m.def("split", [](const random::PRNGKey& key, size_t num) -> std::vector<random::PRNGKey> { return random::Split(key, num); }, "key"_a, "num"_a = 2);

    rand_m.def("fold_in", &random::FoldIn, "key"_a, "data"_a);

    // Array <-> PRNGKey interop for array-backed keys.
    rand_m.def("_key_from_array", [](const Array& a) { return random::PRNGKey(a); }, "array"_a);
    rand_m.def("_key_to_array", [](const random::PRNGKey& k) { return k.ToArray(); }, "key"_a);

    rand_m.def(
        "uniform",
        [](const random::PRNGKey& key, std::vector<int64_t> shape, double minval, double maxval, DType dtype) -> Array {
          return random::Uniform(key, Shape(shape), static_cast<float>(minval), static_cast<float>(maxval), dtype);
        },
        "key"_a, "shape"_a, "minval"_a = 0.0, "maxval"_a = 1.0, "dtype"_a = DType::kFloat32);

    rand_m.def(
        "normal",
        [](const random::PRNGKey& key, std::vector<int64_t> shape, float mean, float stddev, DType dtype) -> Array {
          return random::Normal(key, Shape(shape), mean, stddev, dtype);
        },
        "key"_a, "shape"_a, "mean"_a = 0.0f, "stddev"_a = 1.0f, "dtype"_a = DType::kFloat32);

    rand_m.def(
        "bernoulli",
        [](const random::PRNGKey& key, double p, std::optional<std::vector<int64_t>> shape) -> Array {
          Shape sh = shape.has_value() ? Shape(*shape) : Shape({});
          return random::Bernoulli(key, sh, static_cast<float>(p));
        },
        "key"_a, "p"_a = 0.5, "shape"_a = nb::none());

    rand_m.def(
        "rademacher", [](const random::PRNGKey& key, std::vector<int64_t> shape, DType dtype) -> Array { return random::Rademacher(key, Shape(shape), dtype); },
        "key"_a, "shape"_a, "dtype"_a = DType::kFloat32);

    rand_m.def(
        "categorical",
        [](const random::PRNGKey& key, const Array& logits, std::optional<std::vector<int64_t>> shape) -> Array {
          Shape sh = shape.has_value() ? Shape(*shape) : logits.shape();
          return random::Categorical(key, logits, sh);
        },
        "key"_a, "logits"_a, "shape"_a = nb::none());

    rand_m.def(
        "glorot_uniform",
        [](const random::PRNGKey& key, std::vector<int64_t> shape, int64_t fan_in, int64_t fan_out, DType dtype) -> Array {
          return random::GlorotUniform(key, Shape(shape), fan_in, fan_out, dtype);
        },
        "key"_a, "shape"_a, "fan_in"_a, "fan_out"_a, "dtype"_a = DType::kFloat32);

    rand_m.def(
        "glorot_normal",
        [](const random::PRNGKey& key, std::vector<int64_t> shape, int64_t fan_in, int64_t fan_out, DType dtype) -> Array {
          return random::GlorotNormal(key, Shape(shape), fan_in, fan_out, dtype);
        },
        "key"_a, "shape"_a, "fan_in"_a, "fan_out"_a, "dtype"_a = DType::kFloat32);

    rand_m.def(
        "kaiming_uniform",
        [](const random::PRNGKey& key, std::vector<int64_t> shape, int64_t fan_in, DType dtype) -> Array {
          return random::KaimingUniform(key, Shape(shape), fan_in, dtype);
        },
        "key"_a, "shape"_a, "fan_in"_a, "dtype"_a = DType::kFloat32);

    rand_m.def(
        "kaiming_normal",
        [](const random::PRNGKey& key, std::vector<int64_t> shape, int64_t fan_in, DType dtype) -> Array {
          return random::KaimingNormal(key, Shape(shape), fan_in, dtype);
        },
        "key"_a, "shape"_a, "fan_in"_a, "dtype"_a = DType::kFloat32);
  }

}  // namespace tenzr::python