include/tenzr/numpy/numpy.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_NUMPY_NUMPY_H_
#define TENZR_NUMPY_NUMPY_H_

#include <optional>
#include <string>
#include <utility>
#include <vector>

#include "tenzr/array.h"
#include "tenzr/common/export.h"
#include "tenzr/frontend/dtype.h"
#include "tenzr/frontend/shape.h"

namespace tenzr::numpy {
  struct BroadcastResult {
    Array a;
    Array b;
  };


  // Type Promotion & Broadcasting Engine

  [[nodiscard]] TENZR_API DType PromoteTypes(DType a, DType b);

  [[nodiscard]] TENZR_API Shape BroadcastShapes(const Shape& a, const Shape& b);

  [[nodiscard]] TENZR_API BroadcastResult BroadcastArrays(const Array& a, const Array& b);

  [[nodiscard]] TENZR_API BroadcastResult PromoteAndBroadcast(const Array& a, const Array& b);

  // Array Constructors

  [[nodiscard]] TENZR_API Array Zeros(Shape shape, DType dtype = DType::kFloat32);

  [[nodiscard]] TENZR_API Array Ones(Shape shape, DType dtype = DType::kFloat32);

  [[nodiscard]] TENZR_API Array Full(Shape shape, double fill_value, DType dtype = DType::kFloat32);

  [[nodiscard]] TENZR_API Array Arange(double start, double stop, double step = 1.0, DType dtype = DType::kFloat32);

  [[nodiscard]] TENZR_API Array Arange(double stop, DType dtype = DType::kFloat32);

  [[nodiscard]] TENZR_API Array Linspace(double start, double stop, int64_t num = 50, DType dtype = DType::kFloat32);

  [[nodiscard]] TENZR_API Array Eye(int64_t N, int64_t M = -1, int64_t k = 0, DType dtype = DType::kFloat32);

  template <typename T> [[nodiscard]] inline Array Asarray(const std::vector<T>& data, Shape shape) {
    Array a = Array::Parameter(shape, DtypeOf<T>());
    a.SetData(data);
    if (txpr::TracerContext::IsActive()) {
      auto* tape = txpr::TracerContext::CurrentTape();
      return Array::Symbolic(a.ObtainVarId(tape), std::move(shape), DtypeOf<T>());
    }
    return a;
  }

  template <typename T> [[nodiscard]] inline Array array(const std::vector<T>& data, Shape shape) { return Asarray(data, std::move(shape)); }

  // Binary Mathematical Operations (auto-promotes & broadcasts)

  [[nodiscard]] TENZR_API Array Add(const Array& a, const Array& b);
  [[nodiscard]] TENZR_API Array Subtract(const Array& a, const Array& b);
  [[nodiscard]] TENZR_API Array Multiply(const Array& a, const Array& b);
  [[nodiscard]] TENZR_API Array Divide(const Array& a, const Array& b);
  [[nodiscard]] TENZR_API Array Power(const Array& a, const Array& b);
  [[nodiscard]] TENZR_API Array Maximum(const Array& a, const Array& b);
  [[nodiscard]] TENZR_API Array Minimum(const Array& a, const Array& b);
  [[nodiscard]] TENZR_API Array Remainder(const Array& a, const Array& b);
  [[nodiscard]] TENZR_API Array Atan2(const Array& a, const Array& b);

  // Comparisons (returns DType::Bool array)

  [[nodiscard]] TENZR_API Array Equal(const Array& a, const Array& b);
  [[nodiscard]] TENZR_API Array NotEqual(const Array& a, const Array& b);
  [[nodiscard]] TENZR_API Array Less(const Array& a, const Array& b);
  [[nodiscard]] TENZR_API Array LessEqual(const Array& a, const Array& b);
  [[nodiscard]] TENZR_API Array Greater(const Array& a, const Array& b);
  [[nodiscard]] TENZR_API Array GreaterEqual(const Array& a, const Array& b);

  // Unary Mathematical Functions

  [[nodiscard]] TENZR_API Array Sin(const Array& a);
  [[nodiscard]] TENZR_API Array Cos(const Array& a);
  [[nodiscard]] TENZR_API Array Tan(const Array& a);
  [[nodiscard]] TENZR_API Array Exp(const Array& a);
  [[nodiscard]] TENZR_API Array Log(const Array& a);
  [[nodiscard]] TENZR_API Array Sqrt(const Array& a);
  [[nodiscard]] TENZR_API Array Rsqrt(const Array& a);
  [[nodiscard]] TENZR_API Array Abs(const Array& a);
  [[nodiscard]] TENZR_API Array tanh(const Array& a);
  [[nodiscard]] TENZR_API Array Ceil(const Array& a);
  [[nodiscard]] TENZR_API Array Floor(const Array& a);
  [[nodiscard]] TENZR_API Array Round(const Array& a);
  [[nodiscard]] TENZR_API Array Sign(const Array& a);
  [[nodiscard]] TENZR_API Array Negative(const Array& a);
  [[nodiscard]] TENZR_API Array Clip(const Array& a, const Array& a_min, const Array& a_max);

  // Linear Algebra

  [[nodiscard]] TENZR_API Array Matmul(const Array& a, const Array& b);
  [[nodiscard]] TENZR_API Array Dot(const Array& a, const Array& b);

  // Shape Transformations & View Manipulations

  [[nodiscard]] TENZR_API Array Transpose(const Array& a, const std::vector<int64_t>& axes = {});
  [[nodiscard]] TENZR_API Array Reshape(const Array& a, const Shape& new_shape);
  [[nodiscard]] TENZR_API Array Squeeze(const Array& a, std::optional<int64_t> axis = std::nullopt);
  [[nodiscard]] TENZR_API Array ExpandDims(const Array& a, int64_t axis);
  [[nodiscard]] TENZR_API Array Concatenate(const std::vector<Array>& arrays, int64_t axis = 0);
  [[nodiscard]] TENZR_API Array Stack(const std::vector<Array>& arrays, int64_t axis = 0);

  // Reductions

  [[nodiscard]] TENZR_API Array Sum(const Array& a, std::optional<std::vector<int64_t>> axis = std::nullopt, bool keepdims = false);
  [[nodiscard]] TENZR_API Array Mean(const Array& a, std::optional<std::vector<int64_t>> axis = std::nullopt, bool keepdims = false);
  [[nodiscard]] TENZR_API Array Prod(const Array& a, std::optional<std::vector<int64_t>> axis = std::nullopt, bool keepdims = false);
  [[nodiscard]] TENZR_API Array Max(const Array& a, std::optional<std::vector<int64_t>> axis = std::nullopt, bool keepdims = false);
  [[nodiscard]] TENZR_API Array Min(const Array& a, std::optional<std::vector<int64_t>> axis = std::nullopt, bool keepdims = false);
  [[nodiscard]] TENZR_API Array Var(const Array& a, std::optional<std::vector<int64_t>> axis = std::nullopt, bool keepdims = false);
  [[nodiscard]] TENZR_API Array StdDev(const Array& a, std::optional<std::vector<int64_t>> axis = std::nullopt, bool keepdims = false);

  // Selection & Indexing

  [[nodiscard]] TENZR_API Array Where(const Array& condition, const Array& x, const Array& y);
  [[nodiscard]] TENZR_API Array Take(const Array& a, const Array& indices, int64_t axis = 0);

}  // namespace tenzr::numpy

namespace tenzr {
  namespace tnp = tenzr::numpy;
}  // namespace tenzr

#endif  // TENZR_NUMPY_NUMPY_H_