// 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/numpy.h"
#include <algorithm>
#include <cmath>
#include <cstring>
#include <stdexcept>
#include <string>
#include <utility>
#include <vector>
#include "tenzr/common/error.h"
#include "tenzr/frontend/dtype.h"
#include "tenzr/frontend/memory_pool.h"
#include "tenzr/plax/plax.h"
namespace tenzr::numpy {
// Type Promotion & Broadcasting Engine
DType PromoteTypes(DType a, DType b) { return tenzr::PromoteTypes(a, b); }
Shape BroadcastShapes(const Shape& a, const Shape& b) { return Shape::Broadcast(a, b); }
BroadcastResult BroadcastArrays(const Array& a, const Array& b) {
Shape target = BroadcastShapes(a.shape(), b.shape());
return {a.BroadcastTo(target).Contiguous(), b.BroadcastTo(target).Contiguous()};
}
BroadcastResult PromoteAndBroadcast(const Array& a, const Array& b) {
DType target_dtype = tenzr::PromoteTypes(a.Dtype(), b.Dtype());
Array ca = (a.Dtype() == target_dtype) ? a : a.Astype(target_dtype);
Array cb = (b.Dtype() == target_dtype) ? b : b.Astype(target_dtype);
Shape target_shape = BroadcastShapes(ca.shape(), cb.shape());
return {ca.BroadcastTo(target_shape).Contiguous(), cb.BroadcastTo(target_shape).Contiguous()};
}
// Array Constructors
Array Zeros(Shape shape, DType dtype) { return Array::Zeros(std::move(shape), dtype); }
Array Ones(Shape shape, DType dtype) { return Array::Ones(std::move(shape), dtype); }
Array Full(Shape shape, double fill_value, DType dtype) { return Array::Full(std::move(shape), fill_value, dtype); }
Array Arange(double start, double stop, double step, DType dtype) {
TENZR_CHECK(step != 0.0, "arange step cannot be 0");
std::vector<float> vals;
if (step > 0.0) {
for (double v = start; v < stop; v += step) {
vals.push_back(static_cast<float>(v));
}
} else {
for (double v = start; v > stop; v += step) {
vals.push_back(static_cast<float>(v));
}
}
Array res = Array::Constant(vals, Shape({static_cast<int64_t>(vals.size())}));
if (dtype != DType::kFloat32) {
return res.Astype(dtype);
}
return res;
}
Array Arange(double stop, DType dtype) { return Arange(0.0, stop, 1.0, dtype); }
Array Linspace(double start, double stop, int64_t num, DType dtype) {
TENZR_CHECK(num > 0, "linspace num must be positive");
std::vector<float> vals(static_cast<size_t>(num));
if (num == 1) {
vals[0] = static_cast<float>(start);
} else {
double step = (stop - start) / static_cast<double>(num - 1);
for (int64_t i = 0; i < num; ++i) {
vals[static_cast<size_t>(i)] = static_cast<float>(start + static_cast<double>(i) * step);
}
}
Array res = Array::Constant(vals, Shape({num}));
if (dtype != DType::kFloat32) {
return res.Astype(dtype);
}
return res;
}
Array Eye(int64_t N, int64_t M, int64_t k, DType dtype) {
TENZR_CHECK(N > 0, "eye rows N must be positive");
if (M < 0) {
M = N;
}
TENZR_CHECK(M > 0, "eye cols M must be positive");
std::vector<float> vals(static_cast<size_t>(N * M), 0.0f);
for (int64_t r = 0; r < N; ++r) {
int64_t c = r + k;
if (c >= 0 && c < M) {
vals[static_cast<size_t>(r * M + c)] = 1.0f;
}
}
Array res = Array::Constant(vals, Shape({N, M}));
if (dtype != DType::kFloat32) {
return res.Astype(dtype);
}
return res;
}
// Binary Mathematical Operations (auto-promotes & broadcasts)
Array Add(const Array& a, const Array& b) {
auto [x, y] = PromoteAndBroadcast(a, b);
return plax::Add(x, y);
}
Array Subtract(const Array& a, const Array& b) {
auto [x, y] = PromoteAndBroadcast(a, b);
return plax::Sub(x, y);
}
Array Multiply(const Array& a, const Array& b) {
auto [x, y] = PromoteAndBroadcast(a, b);
return plax::Mul(x, y);
}
Array Divide(const Array& a, const Array& b) {
auto [x, y] = PromoteAndBroadcast(a, b);
return plax::Div(x, y);
}
Array Power(const Array& a, const Array& b) {
auto [x, y] = PromoteAndBroadcast(a, b);
return plax::Pow(x, y);
}
Array Maximum(const Array& a, const Array& b) {
auto [x, y] = PromoteAndBroadcast(a, b);
return plax::Maximum(x, y);
}
Array Minimum(const Array& a, const Array& b) {
auto [x, y] = PromoteAndBroadcast(a, b);
return plax::Minimum(x, y);
}
Array Remainder(const Array& a, const Array& b) {
auto [x, y] = PromoteAndBroadcast(a, b);
return plax::Remainder(x, y);
}
Array Atan2(const Array& a, const Array& b) {
auto [x, y] = PromoteAndBroadcast(a, b);
return plax::Atan2(x, y);
}
// Comparisons (returns DType::Bool array)
Array Equal(const Array& a, const Array& b) {
auto [x, y] = PromoteAndBroadcast(a, b);
return plax::Eq(x, y);
}
Array NotEqual(const Array& a, const Array& b) {
auto [x, y] = PromoteAndBroadcast(a, b);
return plax::Ne(x, y);
}
Array Less(const Array& a, const Array& b) {
auto [x, y] = PromoteAndBroadcast(a, b);
return plax::Lt(x, y);
}
Array LessEqual(const Array& a, const Array& b) {
auto [x, y] = PromoteAndBroadcast(a, b);
return plax::Le(x, y);
}
Array Greater(const Array& a, const Array& b) {
auto [x, y] = PromoteAndBroadcast(a, b);
return plax::Gt(x, y);
}
Array GreaterEqual(const Array& a, const Array& b) {
auto [x, y] = PromoteAndBroadcast(a, b);
return plax::Ge(x, y);
}
// Unary Mathematical Functions
Array Sin(const Array& a) { return plax::Sin(a); }
Array Cos(const Array& a) { return plax::Cos(a); }
Array Tan(const Array& a) { return plax::Tan(a); }
Array Exp(const Array& a) { return plax::Exp(a); }
Array Log(const Array& a) { return plax::Log(a); }
Array Sqrt(const Array& a) { return plax::Sqrt(a); }
Array Rsqrt(const Array& a) { return plax::Rsqrt(a); }
Array Abs(const Array& a) { return plax::Abs(a); }
Array tanh(const Array& a) { return plax::tanh(a); }
Array Ceil(const Array& a) { return plax::Ceil(a); }
Array Floor(const Array& a) { return plax::Floor(a); }
Array Round(const Array& a) { return plax::Round(a); }
Array Sign(const Array& a) { return plax::Sign(a); }
Array Negative(const Array& a) { return plax::Neg(a); }
Array Clip(const Array& a, const Array& a_min, const Array& a_max) {
DType dt = tenzr::PromoteTypes(tenzr::PromoteTypes(a.Dtype(), a_min.Dtype()), a_max.Dtype());
Array ca = (a.Dtype() == dt) ? a : a.Astype(dt);
Array cmin = (a_min.Dtype() == dt) ? a_min : a_min.Astype(dt);
Array cmax = (a_max.Dtype() == dt) ? a_max : a_max.Astype(dt);
Shape s = BroadcastShapes(BroadcastShapes(ca.shape(), cmin.shape()), cmax.shape());
return plax::Clamp(cmin.BroadcastTo(s), ca.BroadcastTo(s), cmax.BroadcastTo(s));
}
// Linear Algebra
Array Matmul(const Array& a, const Array& b) { return a.Matmul(b); }
Array Dot(const Array& a, const Array& b) {
if (a.shape().Rank() == 1 && b.shape().Rank() == 1) {
return Multiply(a, b).Sum();
}
return a.Matmul(b);
}
// Shape Transformations & View Manipulations
Array Transpose(const Array& a, const std::vector<int64_t>& axes) {
if (axes.empty()) {
return a.Transpose();
}
return a.Permute(axes);
}
Array Reshape(const Array& a, const Shape& new_shape) { return a.Reshape(new_shape); }
Array Squeeze(const Array& a, std::optional<int64_t> axis) { return a.Squeeze(axis); }
Array ExpandDims(const Array& a, int64_t axis) { return a.Unsqueeze(axis); }
namespace {
template <typename T> Array concatenate_typed(const std::vector<Array>& arrays, int64_t axis, const Shape& out_shape, DType common_dtype,
int64_t outer_size, int64_t inner_size, int64_t total_axis_dim) {
std::vector<T> out_data(static_cast<size_t>(out_shape.NumElements()), T{});
int64_t cur_axis_offset = 0;
for (const auto& arr : arrays) {
Array casted = (arr.Dtype() == common_dtype) ? arr : arr.Astype(common_dtype);
std::vector<T> in_data = casted.template ToVector<T>();
int64_t arr_axis_dim = casted.shape().Dim(static_cast<size_t>(axis));
for (int64_t o = 0; o < outer_size; ++o) {
int64_t src_offset = o * arr_axis_dim * inner_size;
int64_t dst_offset = (o * total_axis_dim + cur_axis_offset) * inner_size;
if constexpr (std::is_same_v<T, bool>) {
for (int64_t k = 0; k < arr_axis_dim * inner_size; ++k) {
out_data[static_cast<size_t>(dst_offset + k)] = in_data[static_cast<size_t>(src_offset + k)];
}
} else {
std::memcpy(&out_data[static_cast<size_t>(dst_offset)], &in_data[static_cast<size_t>(src_offset)],
static_cast<size_t>(arr_axis_dim * inner_size) * sizeof(T));
}
}
cur_axis_offset += arr_axis_dim;
}
Array res = Array::Zeros(out_shape, common_dtype);
res.SetData(out_data);
return res;
}
} // namespace
Array Concatenate(const std::vector<Array>& arrays, int64_t axis) {
TENZR_CHECK(!arrays.empty(), "numpy::concatenate requires at least one array");
size_t rank = arrays[0].shape().Rank();
if (axis < 0) {
axis += static_cast<int64_t>(rank);
}
TENZR_CHECK(axis >= 0 && axis < static_cast<int64_t>(rank), "numpy::concatenate axis out of range");
// Determine common dtype and total length along axis
DType common_dtype = arrays[0].Dtype();
for (size_t i = 1; i < arrays.size(); ++i) {
common_dtype = tenzr::PromoteTypes(common_dtype, arrays[i].Dtype());
TENZR_CHECK(arrays[i].shape().Rank() == rank, "All arrays in concatenate must have the same rank");
for (size_t d = 0; d < rank; ++d) {
if (static_cast<int64_t>(d) != axis) {
TENZR_CHECK(arrays[i].shape().Dim(d) == arrays[0].shape().Dim(d), "Shape mismatch on non-concatenated axis in concatenate");
}
}
}
int64_t total_axis_dim = 0;
for (const auto& arr : arrays) {
total_axis_dim += arr.shape().Dim(static_cast<size_t>(axis));
}
std::vector<int64_t> out_dims = arrays[0].shape().Dims();
out_dims[static_cast<size_t>(axis)] = total_axis_dim;
Shape out_shape(std::move(out_dims));
int64_t outer_size = 1;
for (int64_t i = 0; i < axis; ++i) {
outer_size *= out_shape.Dim(static_cast<size_t>(i));
}
int64_t inner_size = 1;
for (size_t i = static_cast<size_t>(axis + 1); i < rank; ++i) {
inner_size *= out_shape.Dim(i);
}
switch (common_dtype) {
case DType::kBool:
return concatenate_typed<bool>(arrays, axis, out_shape, common_dtype, outer_size, inner_size, total_axis_dim);
case DType::kInt8:
return concatenate_typed<int8_t>(arrays, axis, out_shape, common_dtype, outer_size, inner_size, total_axis_dim);
case DType::kUInt8:
return concatenate_typed<uint8_t>(arrays, axis, out_shape, common_dtype, outer_size, inner_size, total_axis_dim);
case DType::kInt16:
return concatenate_typed<int16_t>(arrays, axis, out_shape, common_dtype, outer_size, inner_size, total_axis_dim);
case DType::kUInt16:
return concatenate_typed<uint16_t>(arrays, axis, out_shape, common_dtype, outer_size, inner_size, total_axis_dim);
case DType::kInt32:
return concatenate_typed<int32_t>(arrays, axis, out_shape, common_dtype, outer_size, inner_size, total_axis_dim);
case DType::kUInt32:
return concatenate_typed<uint32_t>(arrays, axis, out_shape, common_dtype, outer_size, inner_size, total_axis_dim);
case DType::kInt64:
return concatenate_typed<int64_t>(arrays, axis, out_shape, common_dtype, outer_size, inner_size, total_axis_dim);
case DType::kUInt64:
return concatenate_typed<uint64_t>(arrays, axis, out_shape, common_dtype, outer_size, inner_size, total_axis_dim);
case DType::kFloat16:
return concatenate_typed<tenzr::float16>(arrays, axis, out_shape, common_dtype, outer_size, inner_size, total_axis_dim);
case DType::kBFloat16:
return concatenate_typed<tenzr::bfloat16>(arrays, axis, out_shape, common_dtype, outer_size, inner_size, total_axis_dim);
case DType::kFloat32:
return concatenate_typed<float>(arrays, axis, out_shape, common_dtype, outer_size, inner_size, total_axis_dim);
case DType::kFloat64:
return concatenate_typed<double>(arrays, axis, out_shape, common_dtype, outer_size, inner_size, total_axis_dim);
case DType::kFloat8E4M3FN:
return concatenate_typed<tenzr::float8_e4m3fn>(arrays, axis, out_shape, common_dtype, outer_size, inner_size, total_axis_dim);
case DType::kFloat8E5M2:
return concatenate_typed<tenzr::float8_e5m2>(arrays, axis, out_shape, common_dtype, outer_size, inner_size, total_axis_dim);
case DType::kFloat8E3M4:
return concatenate_typed<tenzr::float8_e3m4>(arrays, axis, out_shape, common_dtype, outer_size, inner_size, total_axis_dim);
case DType::kFloat8E4M3:
return concatenate_typed<tenzr::float8_e4m3>(arrays, axis, out_shape, common_dtype, outer_size, inner_size, total_axis_dim);
case DType::kFloat8E4M3FNUZ:
return concatenate_typed<tenzr::float8_e4m3fnuz>(arrays, axis, out_shape, common_dtype, outer_size, inner_size, total_axis_dim);
case DType::kFloat8E4M3B11FNUZ:
return concatenate_typed<tenzr::float8_e4m3b11fnuz>(arrays, axis, out_shape, common_dtype, outer_size, inner_size, total_axis_dim);
case DType::kFloat8E5M2FNUZ:
return concatenate_typed<tenzr::float8_e5m2fnuz>(arrays, axis, out_shape, common_dtype, outer_size, inner_size, total_axis_dim);
case DType::kFloat8E8M0FNU:
return concatenate_typed<tenzr::float8_e8m0fnu>(arrays, axis, out_shape, common_dtype, outer_size, inner_size, total_axis_dim);
case DType::kFloat6E2M3FN:
return concatenate_typed<tenzr::float6_e2m3fn>(arrays, axis, out_shape, common_dtype, outer_size, inner_size, total_axis_dim);
case DType::kFloat6E3M2FN:
return concatenate_typed<tenzr::float6_e3m2fn>(arrays, axis, out_shape, common_dtype, outer_size, inner_size, total_axis_dim);
case DType::kFloat4E2M1FN:
return concatenate_typed<tenzr::float4_e2m1fn>(arrays, axis, out_shape, common_dtype, outer_size, inner_size, total_axis_dim);
case DType::kInt1:
return concatenate_typed<tenzr::int1>(arrays, axis, out_shape, common_dtype, outer_size, inner_size, total_axis_dim);
case DType::kUInt1:
return concatenate_typed<tenzr::uint1>(arrays, axis, out_shape, common_dtype, outer_size, inner_size, total_axis_dim);
case DType::kInt2:
return concatenate_typed<tenzr::int2>(arrays, axis, out_shape, common_dtype, outer_size, inner_size, total_axis_dim);
case DType::kUInt2:
return concatenate_typed<tenzr::uint2>(arrays, axis, out_shape, common_dtype, outer_size, inner_size, total_axis_dim);
case DType::kInt4:
return concatenate_typed<tenzr::int4>(arrays, axis, out_shape, common_dtype, outer_size, inner_size, total_axis_dim);
case DType::kUInt4:
return concatenate_typed<tenzr::uint4>(arrays, axis, out_shape, common_dtype, outer_size, inner_size, total_axis_dim);
default:
TENZR_THROW("Unsupported dtype in numpy::concatenate: " + std::string(tenzr::ToString(common_dtype)));
}
}
Array Stack(const std::vector<Array>& arrays, int64_t axis) {
TENZR_CHECK(!arrays.empty(), "numpy::stack requires at least one array");
std::vector<Array> expanded;
expanded.reserve(arrays.size());
for (const auto& arr : arrays) {
expanded.push_back(ExpandDims(arr, axis));
}
return Concatenate(expanded, axis);
}
// Reductions
Array Sum(const Array& a, std::optional<std::vector<int64_t>> axis, bool keepdims) { return a.Sum(std::move(axis), keepdims); }
Array Mean(const Array& a, std::optional<std::vector<int64_t>> axis, bool keepdims) { return a.Mean(std::move(axis), keepdims); }
Array Prod(const Array& a, std::optional<std::vector<int64_t>> axis, bool keepdims) { return a.Prod(std::move(axis), keepdims); }
Array Max(const Array& a, std::optional<std::vector<int64_t>> axis, bool keepdims) { return a.Max(std::move(axis), keepdims); }
Array Min(const Array& a, std::optional<std::vector<int64_t>> axis, bool keepdims) { return a.Min(std::move(axis), keepdims); }
Array Var(const Array& a, std::optional<std::vector<int64_t>> axis, bool keepdims) { return a.Var(std::move(axis), keepdims); }
Array StdDev(const Array& a, std::optional<std::vector<int64_t>> axis, bool keepdims) { return a.StdDev(std::move(axis), keepdims); }
// Selection & Indexing
Array Where(const Array& condition, const Array& x, const Array& y) {
auto [px, py] = PromoteAndBroadcast(x, y);
Array pcond = condition.BroadcastTo(px.shape());
return plax::Select(pcond, px, py);
}
Array Take(const Array& a, const Array& indices, int64_t axis) { return plax::Gather(a, indices, axis); }
} // namespace tenzr::numpy