src/python/scipy_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 <limits>

#include "bindings.h"

namespace tenzr::python {

  void BindScipy(nb::module_& m) {
    auto scipy_m = m.def_submodule("scipy", "Scientific computing ecosystem for OpenXLA (scipy-equivalent)");

    // special
    auto special_m = scipy_m.def_submodule("special", "Special mathematical functions");
    special_m.def("erf", &scipy::special::Erf, "x"_a);
    special_m.def("erfc", &scipy::special::Erfc, "x"_a);
    special_m.def("expit", &scipy::special::Expit, "x"_a);
    special_m.def("logit", &scipy::special::Logit, "x"_a);
    special_m.def("logsumexp", &scipy::special::Logsumexp, "x"_a, "axis"_a = std::nullopt, "keepdims"_a = false);
    special_m.def("softmax", &scipy::special::softmax, "x"_a, "axis"_a = -1);
    special_m.def("gamma", &scipy::special::Gamma, "x"_a);
    special_m.def("gammaln", &scipy::special::Gammaln, "x"_a);

    // linalg
    auto linalg_m = scipy_m.def_submodule("linalg", "Linear algebra operations");
    linalg_m.def("det", &scipy::linalg::Det, "A"_a);
    linalg_m.def("inv", &scipy::linalg::Inv, "A"_a);
    linalg_m.def("solve", &scipy::linalg::Solve, "A"_a, "b"_a);
    linalg_m.def("cholesky", &scipy::linalg::Cholesky, "A"_a, "lower"_a = true);

    // optimize
    auto opt_m = scipy_m.def_submodule("optimize", "Numerical optimization solvers");

    nb::class_<scipy::optimize::OptimizeResult>(opt_m, "OptimizeResult")
        .def_prop_ro("x", [](const scipy::optimize::OptimizeResult& r) { return r.x; })
        .def_prop_ro("fun", [](const scipy::optimize::OptimizeResult& r) { return r.fun; })
        .def_prop_ro("nit", [](const scipy::optimize::OptimizeResult& r) { return r.nit; })
        .def_prop_ro("success", [](const scipy::optimize::OptimizeResult& r) { return r.success; })
        .def_prop_ro("message", [](const scipy::optimize::OptimizeResult& r) { return r.message; });

    opt_m.def(
        "minimize_cg",
        [](nb::callable fn, const Array& x0, int64_t max_iter, double tol) {
          TENZR_CHECK(max_iter >= 0 && max_iter <= std::numeric_limits<int>::max(), "max_iter out of range for int");
          auto cpp_fn = [fn](const Array& a) -> Array { return nb::cast<Array>(fn(a)); };
          return scipy::optimize::MinimizeCg(cpp_fn, x0, static_cast<int>(max_iter), tol);
        },
        "fn"_a, "x0"_a, "max_iter"_a = 100, "tol"_a = 1.0e-5);

    opt_m.def(
        "minimize_bfgs",
        [](nb::callable fn, const Array& x0, int64_t max_iter, double tol) {
          TENZR_CHECK(max_iter >= 0 && max_iter <= std::numeric_limits<int>::max(), "max_iter out of range for int");
          auto cpp_fn = [fn](const Array& a) -> Array { return nb::cast<Array>(fn(a)); };
          return scipy::optimize::MinimizeBfgs(cpp_fn, x0, static_cast<int>(max_iter), tol);
        },
        "fn"_a, "x0"_a, "max_iter"_a = 100, "tol"_a = 1.0e-5);
  }

}  // namespace tenzr::python