// 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/scipy/scipy.h"
#include <Eigen/Dense>
#include <algorithm>
#include <cmath>
#include <string>
#include <tuple>
#include <vector>
#include "tenzr/common/error.h"
#include "tenzr/frontend/dtype.h"
#include "tenzr/txpr/grad.h"
namespace tenzr::scipy {
namespace special {
Array Erf(const Array& x) {
// Abramowitz & Stegun 7.1.26 polynomial approximation (error < 1.5e-7)
Array p = Array::Full(x.shape(), 0.3275911, x.Dtype());
Array a1 = Array::Full(x.shape(), 0.254829592, x.Dtype());
Array a2 = Array::Full(x.shape(), -0.284496736, x.Dtype());
Array a3 = Array::Full(x.shape(), 1.421413741, x.Dtype());
Array a4 = Array::Full(x.shape(), -1.453152027, x.Dtype());
Array a5 = Array::Full(x.shape(), 1.061405429, x.Dtype());
Array one = Array::Ones(x.shape(), x.Dtype());
Array sign_x = x.Sign();
Array abs_x = x.Abs();
Array t = one / (one + p * abs_x);
Array poly = t * (a1 + t * (a2 + t * (a3 + t * (a4 + t * a5))));
Array exp_neg_x2 = (-(abs_x * abs_x)).Exp();
Array y = one - poly * exp_neg_x2;
return sign_x * y;
}
Array Erfc(const Array& x) {
Array one = Array::Ones(x.shape(), x.Dtype());
return one - Erf(x);
}
Array Expit(const Array& x) {
Array one = Array::Ones(x.shape(), x.Dtype());
return one / (one + (-x).Exp());
}
Array Logit(const Array& x) {
Array one = Array::Ones(x.shape(), x.Dtype());
return (x / (one - x)).Log();
}
Array Logsumexp(const Array& x, std::optional<std::vector<int64_t>> axis, bool keepdims) {
Array max_val = x.Max(axis, true);
Array exp_shifted = (x - max_val).Exp();
Array sum_exp = exp_shifted.Sum(axis, keepdims);
if (!keepdims) {
max_val = x.Max(axis, false);
}
return max_val + sum_exp.Log();
}
Array softmax(const Array& x, int64_t axis) {
Array max_val = x.Max(std::vector<int64_t>{axis}, true);
Array exp_shifted = (x - max_val).Exp();
Array sum_exp = exp_shifted.Sum(std::vector<int64_t>{axis}, true);
return exp_shifted / sum_exp;
}
Array Gammaln(const Array& x) {
// Ramanujan / Stirling asymptotic expansion for x > 0
Array half = Array::Full(x.shape(), 0.5, x.Dtype());
Array c_stirling = Array::Full(x.shape(), 0.9189385332046727, x.Dtype());
Array c12 = Array::Full(x.shape(), 1.0 / 12.0, x.Dtype());
Array c360 = Array::Full(x.shape(), 1.0 / 360.0, x.Dtype());
Array term1 = c_stirling + (x - half) * x.Log() - x;
Array x_inv = Array::Ones(x.shape(), x.Dtype()) / x;
Array term2 = c12 * x_inv - c360 * x_inv * x_inv * x_inv;
return term1 + term2;
}
Array Gamma(const Array& x) { return Gammaln(x).Exp(); }
} // namespace special
namespace linalg {
Array Det(const Array& A) {
TENZR_CHECK(A.shape().Rank() == 2, "scipy::linalg::det requires a 2-D square matrix");
TENZR_CHECK(A.shape().Dim(0) == A.shape().Dim(1), "scipy::linalg::det requires a square matrix");
int64_t n = A.shape().Dim(0);
if (n == 1) {
return A.Reshape(Shape({}));
}
std::vector<float> a_data = A.ToVector<float>();
Eigen::Map<const Eigen::Matrix<float, Eigen::Dynamic, Eigen::Dynamic, Eigen::RowMajor>> mat(a_data.data(), n, n);
float d = mat.determinant();
return Array::Full(Shape({}), static_cast<double>(d), A.Dtype());
}
Array Inv(const Array& A) {
TENZR_CHECK(A.shape().Rank() == 2, "scipy::linalg::inv requires a 2-D square matrix");
TENZR_CHECK(A.shape().Dim(0) == A.shape().Dim(1), "scipy::linalg::inv requires a square matrix");
int64_t n = A.shape().Dim(0);
std::vector<float> a_data = A.ToVector<float>();
Eigen::Map<const Eigen::Matrix<float, Eigen::Dynamic, Eigen::Dynamic, Eigen::RowMajor>> mat(a_data.data(), n, n);
Eigen::Matrix<float, Eigen::Dynamic, Eigen::Dynamic, Eigen::RowMajor> inv_mat = mat.inverse();
std::vector<float> res_vec(static_cast<size_t>(n * n));
std::memcpy(res_vec.data(), inv_mat.data(), static_cast<size_t>(n * n) * sizeof(float));
Array res = Array::Constant(res_vec, A.shape());
if (A.Dtype() != DType::kFloat32) {
return res.Astype(A.Dtype());
}
return res;
}
Array Solve(const Array& A, const Array& b) {
TENZR_CHECK(A.shape().Rank() == 2, "scipy::linalg::solve requires matrix A to be 2-D");
TENZR_CHECK(A.shape().Dim(0) == A.shape().Dim(1), "scipy::linalg::solve requires matrix A to be square");
TENZR_CHECK(b.shape().Rank() == 1, "scipy::linalg::solve currently requires 1-D vector b");
TENZR_CHECK(A.shape().Dim(0) == b.shape().Dim(0), "scipy::linalg::solve dimension mismatch between A and b");
int64_t n = A.shape().Dim(0);
std::vector<float> a_data = A.ToVector<float>();
std::vector<float> b_data = b.ToVector<float>();
Eigen::Map<const Eigen::Matrix<float, Eigen::Dynamic, Eigen::Dynamic, Eigen::RowMajor>> mat(a_data.data(), n, n);
Eigen::Map<const Eigen::VectorXf> vec(b_data.data(), n);
Eigen::VectorXf x = mat.colPivHouseholderQr().solve(vec);
std::vector<float> res_vec(static_cast<size_t>(n));
std::memcpy(res_vec.data(), x.data(), static_cast<size_t>(n) * sizeof(float));
Array res = Array::Constant(res_vec, b.shape());
if (b.Dtype() != DType::kFloat32) {
return res.Astype(b.Dtype());
}
return res;
}
Array Cholesky(const Array& A, bool lower) {
TENZR_CHECK(A.shape().Rank() == 2, "scipy::linalg::cholesky requires a 2-D square matrix");
TENZR_CHECK(A.shape().Dim(0) == A.shape().Dim(1), "scipy::linalg::cholesky requires a square matrix");
int64_t n = A.shape().Dim(0);
std::vector<float> a_data = A.ToVector<float>();
Eigen::Map<const Eigen::Matrix<float, Eigen::Dynamic, Eigen::Dynamic, Eigen::RowMajor>> mat(a_data.data(), n, n);
Eigen::LLT<Eigen::Matrix<float, Eigen::Dynamic, Eigen::Dynamic, Eigen::RowMajor>> llt(mat);
TENZR_CHECK(llt.info() == Eigen::Success, "scipy::linalg::cholesky failed: matrix is not positive-definite");
Eigen::Matrix<float, Eigen::Dynamic, Eigen::Dynamic, Eigen::RowMajor> res_mat;
if (lower) {
res_mat = llt.matrixL();
} else {
res_mat = llt.matrixU();
}
std::vector<float> res_vec(static_cast<size_t>(n * n));
std::memcpy(res_vec.data(), res_mat.data(), static_cast<size_t>(n * n) * sizeof(float));
Array res = Array::Constant(res_vec, A.shape());
if (A.Dtype() != DType::kFloat32) {
return res.Astype(A.Dtype());
}
return res;
}
} // namespace linalg
namespace optimize {
OptimizeResult MinimizeCg(const std::function<Array(const Array&)>& fn, const Array& x0, int maxiter, double tol) {
auto vg = tenzr::ValueAndGrad(fn);
Array x = x0;
auto [f_val, g] = vg(x);
Array d = -g;
for (int iter = 0; iter < maxiter; ++iter) {
double g_norm = static_cast<double>(std::sqrt((g * g).Sum().Item<float>()));
if (g_norm < tol) {
return {x, f_val, true, iter, "Optimization terminated successfully (tol reached)."};
}
double alpha = 1.0;
double c1 = 1.0e-4;
double g_dot_d = static_cast<double>((g * d).Sum().Item<float>());
if (g_dot_d >= 0) {
d = -g;
g_dot_d = static_cast<double>((g * d).Sum().Item<float>());
}
Array x_new = x + d * alpha;
auto [f_new, g_new] = vg(x_new);
while (static_cast<double>(f_new.Item<float>()) > static_cast<double>(f_val.Item<float>()) + c1 * alpha * g_dot_d && alpha > 1.0e-12) {
alpha *= 0.5;
x_new = x + d * alpha;
{
auto vgr = vg(x_new);
f_new = vgr.value;
g_new = vgr.grad;
}
}
Array y = g_new - g;
float denom = (g * g).Sum().Item<float>();
float beta = denom > 1.0e-12f ? std::max(0.0f, (g_new * y).Sum().Item<float>() / denom) : 0.0f;
d = -g_new + d * static_cast<double>(beta);
x = x_new;
f_val = f_new;
g = g_new;
}
return {x, f_val, false, maxiter, "Maximum iterations reached."};
}
OptimizeResult MinimizeBfgs(const std::function<Array(const Array&)>& fn, const Array& x0, int maxiter, double tol) {
int64_t n = x0.shape().NumElements();
auto vg = tenzr::ValueAndGrad(fn);
Array x = x0;
auto [f_val, g] = vg(x);
// Initialize inverse Hessian approximation H as identity
Array H = Array::Full(Shape({n, n}), 0.0, x0.Dtype());
std::vector<float> h_eye(static_cast<size_t>(n * n), 0.0f);
for (int64_t i = 0; i < n; ++i) {
h_eye[static_cast<size_t>(i * n + i)] = 1.0f;
}
H.SetData(h_eye);
for (int iter = 0; iter < maxiter; ++iter) {
double g_norm = static_cast<double>(std::sqrt((g * g).Sum().Item<float>()));
if (g_norm < tol) {
return {x, f_val, true, iter, "Optimization terminated successfully (tol reached)."};
}
// Search direction d = -H * g
Array d = -H.Matmul(g.Reshape(Shape({n, 1}))).Reshape(x0.shape());
double alpha = 1.0;
double c1 = 1.0e-4;
double g_dot_d = static_cast<double>((g * d).Sum().Item<float>());
if (g_dot_d >= 0) {
d = -g;
g_dot_d = static_cast<double>((g * d).Sum().Item<float>());
}
Array x_new = x + d * alpha;
auto [f_new, g_new] = vg(x_new);
while (static_cast<double>(f_new.Item<float>()) > static_cast<double>(f_val.Item<float>()) + c1 * alpha * g_dot_d && alpha > 1.0e-12) {
alpha *= 0.5;
x_new = x + d * alpha;
{
auto vgr = vg(x_new);
f_new = vgr.value;
g_new = vgr.grad;
}
}
Array s = (x_new - x).Reshape(Shape({n, 1}));
Array y = (g_new - g).Reshape(Shape({n, 1}));
float y_dot_s = (y * s).Sum().Item<float>();
if (y_dot_s > 1.0e-12f) {
float rho = 1.0f / y_dot_s;
Array I = Array::Full(Shape({n, n}), 0.0, x0.Dtype());
I.SetData(h_eye);
Array s_yT = s.Matmul(y.Transpose());
Array y_sT = y.Matmul(s.Transpose());
Array s_sT = s.Matmul(s.Transpose());
Array V1 = I - s_yT * static_cast<double>(rho);
Array V2 = I - y_sT * static_cast<double>(rho);
H = V1.Matmul(H).Matmul(V2) + s_sT * static_cast<double>(rho);
}
x = x_new;
f_val = f_new;
g = g_new;
}
return {x, f_val, false, maxiter, "Maximum iterations reached."};
}
} // namespace optimize
} // namespace tenzr::scipy