include/tenzr/frontend/array.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_FRONTEND_ARRAY_H_
#define TENZR_FRONTEND_ARRAY_H_

#include <cstring>
#include <initializer_list>
#include <memory>
#include <string>
#include <vector>

#include "tenzr/common/error.h"
#include "tenzr/frontend/dtype.h"
#include "tenzr/frontend/layout.h"
#include "tenzr/frontend/memory_pool.h"
#include "tenzr/frontend/node.h"
#include "tenzr/frontend/shape.h"
#include "tenzr/frontend/sharding.h"
#include "tenzr/txpr/tracer.h"
#include "tenzr/txpr/txpr.h"

namespace tenzr {

  class Array {
   public:
    Array() = default;

    explicit Array(std::shared_ptr<Node> node);

    // Factory methods

    static Array Parameter(Shape shape, DType dtype = DType::kFloat32, std::string name = "");

    static Array Constant(const std::vector<float>& data, Shape shape);

    static Array Zeros(Shape shape, DType dtype = DType::kFloat32);

    static Array Ones(Shape shape, DType dtype = DType::kFloat32);

    static Array Full(Shape shape, double value, DType dtype = DType::kFloat32);

    // Metadata accessors

    const Shape& shape() const;

    DType Dtype() const;

    const ShardingSpec& Sharding() const;

    void Shard(ShardingSpec spec);

    void Shard(const NamedSharding& ns);

    const StridedLayout& Layout() const;

    bool IsContiguous() const;

    const std::shared_ptr<Node>& node() const { return node_; }

    bool IsDefined() const { return node_ != nullptr || var_id_.IsValid(); }

    // Dynamic tracing / symbolic handles

    [[nodiscard]] bool IsSymbolic() const noexcept { return var_id_.IsValid(); }

    [[nodiscard]] txpr::VarId SymbolicVarId() const noexcept { return var_id_; }

    [[nodiscard]] uint64_t Id() const noexcept {
      if (var_id_.IsValid()) {
        return static_cast<uint64_t>(var_id_.id);
      }
      return node_ ? node_->Id() : 0;
    }

    static Array Symbolic(txpr::VarId vid, Shape shape, DType dtype) {
      Array a;
      a.var_id_ = vid;
      a.shape_ = std::move(shape);
      a.dtype_ = dtype;
      return a;
    }

    [[nodiscard]] txpr::VarId ObtainVarId(txpr::Txpr* tape) const;

    // Host data management

    void BindData(std::shared_ptr<PinnedBuffer> buffer);

    std::shared_ptr<PinnedBuffer> Data() const;

    bool HasData() const;

    bool IsLazy() const { return !HasData(); }

    // Device data management

    void BindDeviceData(std::shared_ptr<DeviceBuffer> buffer);

    std::shared_ptr<DeviceBuffer> DeviceData() const;

    bool HasDeviceData() const;

    // Synchronization

    void SyncHost() const;

    void SyncDevice() const;

    Array Clone() const;

    Array Eval() const;

    Array Materialize() const;

    template <typename T> T Item() const {
      TENZR_CHECK(shape().NumElements() == 1, "item() can only be called on arrays with 1 element");
      return ToVector<T>()[0];
    }

    template <typename T> void SetData(const std::vector<T>& values) {
      TENZR_CHECK(shape().NumElements() == static_cast<int64_t>(values.size()), "Data size does not match array shape");
      TENZR_CHECK(DtypeOf<T>() == Dtype(), "set_data<T>() dtype mismatch: values have dtype " + std::string(tenzr::ToString(DtypeOf<T>())) +
                                               " but array has dtype " + std::string(tenzr::ToString(Dtype())));
      if constexpr (std::is_same_v<T, bool>) {
        size_t bytes = values.size();
        auto buf = MemoryPool::Instance().Allocate(bytes);
        uint8_t* dst = reinterpret_cast<uint8_t*>(buf->Data());
        for (size_t i = 0; i < values.size(); ++i) {
          dst[i] = values[i] ? 1 : 0;
        }
        BindData(buf);
      } else {
        size_t bytes = values.size() * sizeof(T);
        auto buf = MemoryPool::Instance().Allocate(bytes);
        std::memcpy(buf->Data(), values.data(), bytes);
        BindData(buf);
      }
    }

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

    template <typename T> std::vector<T> ToVector() const {
      if (!HasData()) {
        return Materialize().ToVector<T>();
      }
      TENZR_CHECK(DtypeOf<T>() == Dtype(), "to_vector<T>() dtype mismatch: requested " + std::string(tenzr::ToString(DtypeOf<T>())) + " but array has dtype " +
                                               std::string(tenzr::ToString(Dtype())));
      auto buf = Data();
      TENZR_CHECK(buf != nullptr, "Array has no bound data");
      size_t count = static_cast<size_t>(shape().NumElements());

      if constexpr (std::is_same_v<T, bool>) {
        std::vector<bool> result(count);
        const uint8_t* src = reinterpret_cast<const uint8_t*>(buf->Data());
        for (size_t i = 0; i < count; ++i) {
          result[i] = (src[i] != 0);
        }
        return result;
      } else {
        std::vector<T> result(count);
        const T* src = reinterpret_cast<const T*>(buf->Data());

        const auto& lay = Layout();
        if (lay.IsContiguous() && lay.Offset() == 0) {
          std::memcpy(result.data(), src, count * sizeof(T));
        } else {
          size_t r = shape().Rank();
          if (r == 0) {
            result[0] = src[lay.Offset()];
          } else {
            std::vector<int64_t> idx(r, 0);
            for (size_t i = 0; i < count; ++i) {
              int64_t off = lay.LinearOffset(idx);
              result[i] = src[off];
              for (size_t d = r; d > 0; --d) {
                size_t dim_idx = d - 1;
                idx[dim_idx]++;
                if (idx[dim_idx] < shape().Dim(dim_idx) || dim_idx == 0) {
                  break;
                }
                idx[dim_idx] = 0;
              }
            }
          }
        }
        return result;
      }
    }

    Array Astype(DType target_dtype) const;

    template <typename T> std::vector<T> ToVectorAs() const { return Astype(DtypeOf<T>()).template ToVector<T>(); }

    // Registry & Custom Ops

    static Array CustomOp(std::string op_name, std::vector<Array> inputs);

    Array CustomOp(std::string op_name) const;

    // View and layout transformations

    Array Reshape(Shape new_shape) const;

    Array Flatten() const;

    Array Ravel() const;

    Array Transpose(int64_t dim0 = -1, int64_t dim1 = -2) const;

    Array Permute(const std::vector<int64_t>& axes) const;

    Array Squeeze(std::optional<int64_t> dim = std::nullopt) const;

    Array Unsqueeze(int64_t dim) const;

    Array slice(const std::vector<Slice>& slices) const;

    Array BroadcastTo(const Shape& target_shape) const;

    Array AsStrided(const Shape& shape, const std::vector<int64_t>& strides, int64_t offset = 0) const;

    Array Contiguous() const;

    Array Matmul(const Array& other) const;

    // Tier 1: Unary Elementwise Operations

    Array Relu() const;

    Array sigmoid() const;

    Array tanh() const;

    Array Exp() const;

    Array Log() const;

    Array Abs() const;

    Array Neg() const;

    Array Sin() const;

    Array Cos() const;

    Array Sqrt() const;

    Array Rsqrt() const;

    Array Ceil() const;

    Array Floor() const;

    Array Round() const;

    Array Sign() const;

    Array Cbrt() const;

    Array Tan() const;

    Array Expm1() const;

    Array Log1p() const;

    Array IsFinite() const;

    Array LogicalNot() const;

    // Composite activations

    Array Silu() const;

    Array Gelu() const;

    // Tier 2: Binary Elementwise Operations

    Array Pow(const Array& other) const;

    Array Remainder(const Array& other) const;

    Array Atan2(const Array& other) const;

    Array Minimum(const Array& other) const;

    Array Maximum(const Array& other) const;

    Array BitwiseAnd(const Array& other) const;

    Array BitwiseOr(const Array& other) const;

    Array BitwiseXor(const Array& other) const;

    Array ShiftLeft(const Array& other) const;

    Array ShiftRight(const Array& other) const;

    // Tier 3: Clamp & Comparison Operations

    Array Clamp(const Array& min, const Array& max) const;

    Array Equal(const Array& other) const;

    Array NotEqual(const Array& other) const;

    Array GreaterEqual(const Array& other) const;

    Array GreaterThan(const Array& other) const;

    Array LessEqual(const Array& other) const;

    Array LessThan(const Array& other) const;

    // Reductions

    Array Sum(std::optional<std::vector<int64_t>> dims = std::nullopt, bool keepdims = false) const;

    Array Mean(std::optional<std::vector<int64_t>> dims = std::nullopt, bool keepdims = false) const;

    Array Max(std::optional<std::vector<int64_t>> dims = std::nullopt, bool keepdims = false) const;

    Array Min(std::optional<std::vector<int64_t>> dims = std::nullopt, bool keepdims = false) const;

    Array Prod(std::optional<std::vector<int64_t>> dims = std::nullopt, bool keepdims = false) const;

    Array Var(std::optional<std::vector<int64_t>> dims = std::nullopt, bool keepdims = false) const;

    Array StdDev(std::optional<std::vector<int64_t>> dims = std::nullopt, bool keepdims = false) const;

    // Overloaded arithmetic operators

    Array operator+(const Array& other) const;

    Array operator*(const Array& other) const;

    Array operator-(const Array& other) const;

    Array operator/(const Array& other) const;

    // Overloaded comparison operators

    Array operator==(const Array& other) const;

    Array operator!=(const Array& other) const;

    Array operator<(const Array& other) const;

    Array operator<=(const Array& other) const;

    Array operator>(const Array& other) const;

    Array operator>=(const Array& other) const;

    std::string ToString() const;

   private:
    std::shared_ptr<Node> node_;
    txpr::VarId var_id_{txpr::VarId::Invalid()};
    Shape shape_;
    DType dtype_{DType::kFloat32};
  };

  // Factory functions

  inline Array array(const std::vector<float>& data) { return Array::Constant(data, Shape({static_cast<int64_t>(data.size())})); }

  inline Array array(const std::vector<float>& data, Shape shape) { return Array::Constant(data, std::move(shape)); }

  inline Array array(std::initializer_list<float> data) { return Array::Constant(std::vector<float>(data), Shape({static_cast<int64_t>(data.size())})); }

  inline Array array(std::initializer_list<float> data, Shape shape) { return Array::Constant(std::vector<float>(data), std::move(shape)); }

  inline Array array(std::initializer_list<std::initializer_list<float>> matrix_data) {
    int64_t rows = static_cast<int64_t>(matrix_data.size());
    int64_t cols = rows > 0 ? static_cast<int64_t>(matrix_data.begin()->size()) : 0;
    std::vector<float> data;
    data.reserve(static_cast<size_t>(rows * cols));
    for (const auto& row : matrix_data) {
      TENZR_CHECK(static_cast<int64_t>(row.size()) == cols, "All rows in matrix initializer list must have equal column length");
      data.insert(data.end(), row.begin(), row.end());
    }
    return Array::Constant(data, Shape({rows, cols}));
  }

  inline Array Matrix(std::initializer_list<std::initializer_list<float>> m) { return array(m); }

  inline Array Zeros(Shape shape, DType dtype = DType::kFloat32) { return Array::Zeros(std::move(shape), dtype); }

  inline Array Ones(Shape shape, DType dtype = DType::kFloat32) { return Array::Ones(std::move(shape), dtype); }

  inline Array Full(Shape shape, double value, DType dtype = DType::kFloat32) { return Array::Full(std::move(shape), value, dtype); }

  inline Array Sum(const Array& a, std::optional<std::vector<int64_t>> dims = std::nullopt, bool keepdims = false) { return a.Sum(std::move(dims), keepdims); }

  inline Array Mean(const Array& a, std::optional<std::vector<int64_t>> dims = std::nullopt, bool keepdims = false) {
    return a.Mean(std::move(dims), keepdims);
  }

  inline Array Max(const Array& a, std::optional<std::vector<int64_t>> dims = std::nullopt, bool keepdims = false) { return a.Max(std::move(dims), keepdims); }

  inline Array Min(const Array& a, std::optional<std::vector<int64_t>> dims = std::nullopt, bool keepdims = false) { return a.Min(std::move(dims), keepdims); }

  inline Array Prod(const Array& a, std::optional<std::vector<int64_t>> dims = std::nullopt, bool keepdims = false) {
    return a.Prod(std::move(dims), keepdims);
  }

  inline Array Arange(float start, float stop, float step = 1.0f) {
    TENZR_CHECK(step != 0.0f, "arange step cannot be 0");
    std::vector<float> vals;
    if (step > 0.0f) {
      for (float v = start; v < stop; v += step) {
        vals.push_back(v);
      }
    } else {
      for (float v = start; v > stop; v += step) {
        vals.push_back(v);
      }
    }
    return Array::Constant(vals, Shape({static_cast<int64_t>(vals.size())}));
  }

  inline Array Eye(int64_t n, DType dtype = DType::kFloat32) {
    TENZR_CHECK(n > 0, "eye size must be positive");
    TENZR_CHECK(dtype == DType::kFloat32, "eye currently supports kFloat32");
    std::vector<float> vals(static_cast<size_t>(n * n), 0.0f);
    for (int64_t i = 0; i < n; ++i) {
      vals[static_cast<size_t>(i * n + i)] = 1.0f;
    }
    return Array::Constant(vals, Shape({n, n}));
  }

  // Free function mathematical and functional helpers
  inline Array Divide(const Array& a, const Array& b) { return a / b; }

  inline Array Matmul(const Array& a, const Array& b) { return a.Matmul(b); }

  inline Array Reshape(const Array& a, Shape new_shape) { return a.Reshape(std::move(new_shape)); }

  inline Array Relu(const Array& a) { return a.Relu(); }

  inline Array sigmoid(const Array& a) { return a.sigmoid(); }

  inline Array tanh(const Array& a) { return a.tanh(); }

  inline Array Exp(const Array& a) { return a.Exp(); }

  inline Array Log(const Array& a) { return a.Log(); }

  inline Array Abs(const Array& a) { return a.Abs(); }

  inline Array Neg(const Array& a) { return a.Neg(); }

  inline Array Sin(const Array& a) { return a.Sin(); }

  inline Array Cos(const Array& a) { return a.Cos(); }

  inline Array Sqrt(const Array& a) { return a.Sqrt(); }

  inline Array Rsqrt(const Array& a) { return a.Rsqrt(); }

  inline Array Ceil(const Array& a) { return a.Ceil(); }

  inline Array Floor(const Array& a) { return a.Floor(); }

  inline Array Round(const Array& a) { return a.Round(); }

  inline Array Sign(const Array& a) { return a.Sign(); }

  inline Array Cbrt(const Array& a) { return a.Cbrt(); }

  inline Array Tan(const Array& a) { return a.Tan(); }

  inline Array Expm1(const Array& a) { return a.Expm1(); }

  inline Array Log1p(const Array& a) { return a.Log1p(); }

  inline Array IsFinite(const Array& a) { return a.IsFinite(); }

  inline Array LogicalNot(const Array& a) { return a.LogicalNot(); }

  inline Array Silu(const Array& a) { return a.Silu(); }

  inline Array Gelu(const Array& a) { return a.Gelu(); }

  inline Array Pow(const Array& a, const Array& b) { return a.Pow(b); }

  inline Array Remainder(const Array& a, const Array& b) { return a.Remainder(b); }

  inline Array Atan2(const Array& a, const Array& b) { return a.Atan2(b); }

  inline Array Minimum(const Array& a, const Array& b) { return a.Minimum(b); }

  inline Array Maximum(const Array& a, const Array& b) { return a.Maximum(b); }

  inline Array Clamp(const Array& x, const Array& min, const Array& max) { return x.Clamp(min, max); }

  inline Array Select(const Array& pred, const Array& on_true, const Array& on_false) { return Array::CustomOp("select", {pred, on_true, on_false}); }

  inline Array Equal(const Array& a, const Array& b) { return a.Equal(b); }

  inline Array NotEqual(const Array& a, const Array& b) { return a.NotEqual(b); }

  inline Array GreaterEqual(const Array& a, const Array& b) { return a.GreaterEqual(b); }

  inline Array GreaterThan(const Array& a, const Array& b) { return a.GreaterThan(b); }

  inline Array LessEqual(const Array& a, const Array& b) { return a.LessEqual(b); }

  inline Array LessThan(const Array& a, const Array& b) { return a.LessThan(b); }

  // Unary negation operator
  inline Array operator-(const Array& a) { return a.Neg(); }

  // Scalar arithmetic operators (automatic broadcast & promotion)
  inline Array operator+(const Array& a, double s) { return a + Array::Full(a.shape(), s, a.Dtype()); }
  inline Array operator+(double s, const Array& a) { return Array::Full(a.shape(), s, a.Dtype()) + a; }
  inline Array operator-(const Array& a, double s) { return a - Array::Full(a.shape(), s, a.Dtype()); }
  inline Array operator-(double s, const Array& a) { return Array::Full(a.shape(), s, a.Dtype()) - a; }
  inline Array operator*(const Array& a, double s) { return a * Array::Full(a.shape(), s, a.Dtype()); }
  inline Array operator*(double s, const Array& a) { return Array::Full(a.shape(), s, a.Dtype()) * a; }
  inline Array operator/(const Array& a, double s) { return a / Array::Full(a.shape(), s, a.Dtype()); }
  inline Array operator/(double s, const Array& a) { return Array::Full(a.shape(), s, a.Dtype()) / a; }

  std::ostream& operator<<(std::ostream& os, const Array& t);

}  // namespace tenzr

#endif  // TENZR_FRONTEND_ARRAY_H_