// 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/frontend/array.h"
#include <algorithm>
#include <numeric>
#include <sstream>
#include "tenzr/common/error.h"
#include "tenzr/execution/execution_engine.h"
#include "tenzr/frontend/op_registry.h"
#include "tenzr/txpr/custom_vjp.h"
#include "tenzr/txpr/tracer.h"
#include "tenzr/txpr/txpr.h"
namespace tenzr {
Array::Array(std::shared_ptr<Node> node) : node_(std::move(node)) {}
Array Array::Parameter(Shape shape, DType dtype, std::string name) {
auto node = std::make_shared<VariableNode>(std::move(shape), dtype, std::move(name));
return Array(std::move(node));
}
Array Array::Constant(const std::vector<float>& data, Shape shape) {
if (txpr::TracerContext::IsActive() && data.size() == 1) {
return Full(std::move(shape), static_cast<double>(data[0]), DType::kFloat32);
}
auto node = std::make_shared<VariableNode>(shape, DType::kFloat32);
size_t bytes = data.size() * sizeof(float);
auto buf = MemoryPool::Instance().Allocate(bytes);
std::memcpy(buf->Data(), data.data(), bytes);
node->SetData(buf);
return Array(std::move(node));
}
namespace {
template <typename T> void FillBufferTyped(const std::shared_ptr<PinnedBuffer>& buf, size_t count, double value) {
if constexpr (std::is_same_v<T, bool>) {
uint8_t* pb = reinterpret_cast<uint8_t*>(buf->Data());
uint8_t v = (value != 0.0) ? 1 : 0;
for (size_t i = 0; i < count; ++i) {
pb[i] = v;
}
} else {
T* p = buf->As<T>();
T v = static_cast<T>(value);
for (size_t i = 0; i < count; ++i) {
p[i] = v;
}
}
}
} // namespace
Array Array::Full(Shape shape, double value, DType dtype) {
if (txpr::TracerContext::IsActive()) {
auto* tape = txpr::TracerContext::CurrentTape();
if (tape) {
txpr::AttributePayload attrs;
attrs.scalar_value = value;
txpr::VarId out_v = tape->AppendInstruction(txpr::OpCode::kConstant, {}, shape, dtype, std::move(attrs));
return Array::Symbolic(out_v, std::move(shape), dtype);
}
}
auto t = Parameter(shape, dtype);
size_t count = static_cast<size_t>(shape.NumElements());
size_t bytes = count * SizeOf(dtype);
auto buf = MemoryPool::Instance().Allocate(bytes);
switch (dtype) {
case DType::kBool:
FillBufferTyped<bool>(buf, count, value);
break;
case DType::kInt8:
FillBufferTyped<int8_t>(buf, count, value);
break;
case DType::kUInt8:
FillBufferTyped<uint8_t>(buf, count, value);
break;
case DType::kInt16:
FillBufferTyped<int16_t>(buf, count, value);
break;
case DType::kUInt16:
FillBufferTyped<uint16_t>(buf, count, value);
break;
case DType::kInt32:
FillBufferTyped<int32_t>(buf, count, value);
break;
case DType::kUInt32:
FillBufferTyped<uint32_t>(buf, count, value);
break;
case DType::kInt64:
FillBufferTyped<int64_t>(buf, count, value);
break;
case DType::kUInt64:
FillBufferTyped<uint64_t>(buf, count, value);
break;
case DType::kFloat16:
FillBufferTyped<tenzr::float16>(buf, count, value);
break;
case DType::kBFloat16:
FillBufferTyped<tenzr::bfloat16>(buf, count, value);
break;
case DType::kFloat32:
FillBufferTyped<float>(buf, count, value);
break;
case DType::kFloat64:
FillBufferTyped<double>(buf, count, value);
break;
case DType::kFloat8E4M3FN:
FillBufferTyped<tenzr::float8_e4m3fn>(buf, count, value);
break;
case DType::kFloat8E5M2:
FillBufferTyped<tenzr::float8_e5m2>(buf, count, value);
break;
case DType::kFloat8E3M4:
FillBufferTyped<tenzr::float8_e3m4>(buf, count, value);
break;
case DType::kFloat8E4M3:
FillBufferTyped<tenzr::float8_e4m3>(buf, count, value);
break;
case DType::kFloat8E4M3FNUZ:
FillBufferTyped<tenzr::float8_e4m3fnuz>(buf, count, value);
break;
case DType::kFloat8E4M3B11FNUZ:
FillBufferTyped<tenzr::float8_e4m3b11fnuz>(buf, count, value);
break;
case DType::kFloat8E5M2FNUZ:
FillBufferTyped<tenzr::float8_e5m2fnuz>(buf, count, value);
break;
case DType::kFloat8E8M0FNU:
FillBufferTyped<tenzr::float8_e8m0fnu>(buf, count, value);
break;
case DType::kFloat6E2M3FN:
FillBufferTyped<tenzr::float6_e2m3fn>(buf, count, value);
break;
case DType::kFloat6E3M2FN:
FillBufferTyped<tenzr::float6_e3m2fn>(buf, count, value);
break;
case DType::kFloat4E2M1FN:
FillBufferTyped<tenzr::float4_e2m1fn>(buf, count, value);
break;
case DType::kInt1:
FillBufferTyped<tenzr::int1>(buf, count, value);
break;
case DType::kUInt1:
FillBufferTyped<tenzr::uint1>(buf, count, value);
break;
case DType::kInt2:
FillBufferTyped<tenzr::int2>(buf, count, value);
break;
case DType::kUInt2:
FillBufferTyped<tenzr::uint2>(buf, count, value);
break;
case DType::kInt4:
FillBufferTyped<tenzr::int4>(buf, count, value);
break;
case DType::kUInt4:
FillBufferTyped<tenzr::uint4>(buf, count, value);
break;
default:
TENZR_THROW("Unsupported dtype for full: " + std::string(tenzr::ToString(dtype)));
}
t.BindData(buf);
return t;
}
Array Array::Zeros(Shape shape, DType dtype) {
if (txpr::TracerContext::IsActive()) {
return Full(std::move(shape), 0.0, dtype);
}
auto t = Parameter(shape, dtype);
size_t bytes = static_cast<size_t>(shape.NumElements()) * SizeOf(dtype);
auto buf = MemoryPool::Instance().Allocate(bytes);
std::memset(buf->Data(), 0, bytes);
t.BindData(buf);
return t;
}
Array Array::Ones(Shape shape, DType dtype) { return Full(std::move(shape), 1.0, dtype); }
const Shape& Array::shape() const {
if (var_id_.IsValid()) {
return shape_;
}
TENZR_CHECK(node_ != nullptr, "Array is uninitialized");
return node_->shape();
}
DType Array::Dtype() const {
if (var_id_.IsValid()) {
return dtype_;
}
TENZR_CHECK(node_ != nullptr, "Array is uninitialized");
return node_->Dtype();
}
const ShardingSpec& Array::Sharding() const {
if (var_id_.IsValid()) {
static const ShardingSpec default_sharding;
return default_sharding;
}
TENZR_CHECK(node_ != nullptr, "Array is uninitialized");
return node_->Sharding();
}
void Array::Shard(ShardingSpec spec) {
if (var_id_.IsValid()) {
return;
}
TENZR_CHECK(node_ != nullptr, "Array is uninitialized");
node_->SetSharding(std::move(spec));
}
void Array::Shard(const NamedSharding& ns) { Shard(ShardingSpec(ns)); }
void Array::BindData(std::shared_ptr<PinnedBuffer> buffer) {
TENZR_CHECK(node_ != nullptr, "Array is uninitialized");
auto var_node = std::dynamic_pointer_cast<VariableNode>(node_);
TENZR_CHECK(var_node != nullptr, "Cannot bind data to non-variable AST node");
var_node->SetData(std::move(buffer));
}
std::shared_ptr<PinnedBuffer> Array::Data() const {
if (!node_) {
return nullptr;
}
auto var_node = std::dynamic_pointer_cast<VariableNode>(node_);
if (!var_node) {
return nullptr;
}
return var_node->Data();
}
bool Array::HasData() const {
if (var_id_.IsValid()) {
return false;
}
if (!node_) {
return false;
}
auto var_node = std::dynamic_pointer_cast<VariableNode>(node_);
return var_node && var_node->HasData();
}
void Array::BindDeviceData(std::shared_ptr<DeviceBuffer> buffer) {
TENZR_CHECK(node_ != nullptr, "Array is uninitialized");
auto var_node = std::dynamic_pointer_cast<VariableNode>(node_);
TENZR_CHECK(var_node != nullptr, "Cannot bind device data to non-variable AST node");
var_node->SetDeviceData(std::move(buffer));
}
std::shared_ptr<DeviceBuffer> Array::DeviceData() const {
if (!node_) {
return nullptr;
}
auto var_node = std::dynamic_pointer_cast<VariableNode>(node_);
if (!var_node) {
return nullptr;
}
return var_node->DeviceData();
}
bool Array::HasDeviceData() const {
if (var_id_.IsValid()) {
return false;
}
if (!node_) {
return false;
}
auto var_node = std::dynamic_pointer_cast<VariableNode>(node_);
return var_node && var_node->HasDeviceData();
}
txpr::VarId Array::ObtainVarId(txpr::Txpr* tape) const {
if (var_id_.IsValid()) {
return var_id_;
}
TENZR_CHECK(tape != nullptr, "Active tape required to obtain VarId for Array");
if (HasData()) {
if (shape().NumElements() == 1) {
double val = 0.0;
switch (Dtype()) {
case DType::kFloat32:
val = static_cast<double>(Item<float>());
break;
case DType::kFloat64:
val = Item<double>();
break;
case DType::kFloat16:
val = static_cast<double>(Item<tenzr::float16>());
break;
case DType::kBFloat16:
val = static_cast<double>(Item<tenzr::bfloat16>());
break;
case DType::kFloat8E3M4:
val = static_cast<double>(Item<tenzr::float8_e3m4>());
break;
case DType::kFloat8E4M3:
val = static_cast<double>(Item<tenzr::float8_e4m3>());
break;
case DType::kFloat8E4M3FN:
val = static_cast<double>(Item<tenzr::float8_e4m3fn>());
break;
case DType::kFloat8E4M3FNUZ:
val = static_cast<double>(Item<tenzr::float8_e4m3fnuz>());
break;
case DType::kFloat8E4M3B11FNUZ:
val = static_cast<double>(Item<tenzr::float8_e4m3b11fnuz>());
break;
case DType::kFloat8E5M2:
val = static_cast<double>(Item<tenzr::float8_e5m2>());
break;
case DType::kFloat8E5M2FNUZ:
val = static_cast<double>(Item<tenzr::float8_e5m2fnuz>());
break;
case DType::kFloat8E8M0FNU:
val = static_cast<double>(Item<tenzr::float8_e8m0fnu>());
break;
case DType::kFloat6E2M3FN:
val = static_cast<double>(Item<tenzr::float6_e2m3fn>());
break;
case DType::kFloat6E3M2FN:
val = static_cast<double>(Item<tenzr::float6_e3m2fn>());
break;
case DType::kFloat4E2M1FN:
val = static_cast<double>(Item<tenzr::float4_e2m1fn>());
break;
case DType::kInt8:
val = static_cast<double>(Item<int8_t>());
break;
case DType::kUInt8:
val = static_cast<double>(Item<uint8_t>());
break;
case DType::kInt16:
val = static_cast<double>(Item<int16_t>());
break;
case DType::kUInt16:
val = static_cast<double>(Item<uint16_t>());
break;
case DType::kInt32:
val = static_cast<double>(Item<int32_t>());
break;
case DType::kUInt32:
val = static_cast<double>(Item<uint32_t>());
break;
case DType::kInt64:
val = static_cast<double>(Item<int64_t>());
break;
case DType::kUInt64:
val = static_cast<double>(Item<uint64_t>());
break;
case DType::kInt1:
val = static_cast<double>(static_cast<int64_t>(Item<tenzr::int1>()));
break;
case DType::kUInt1:
val = static_cast<double>(static_cast<int64_t>(Item<tenzr::uint1>()));
break;
case DType::kInt2:
val = static_cast<double>(static_cast<int64_t>(Item<tenzr::int2>()));
break;
case DType::kUInt2:
val = static_cast<double>(static_cast<int64_t>(Item<tenzr::uint2>()));
break;
case DType::kInt4:
val = static_cast<double>(static_cast<int64_t>(Item<tenzr::int4>()));
break;
case DType::kUInt4:
val = static_cast<double>(static_cast<int64_t>(Item<tenzr::uint4>()));
break;
case DType::kBool:
val = Item<bool>() ? 1.0 : 0.0;
break;
default:
val = 0.0;
break;
}
txpr::AttributePayload attrs;
attrs.scalar_value = val;
return tape->AppendInstruction(txpr::OpCode::kConstant, {}, shape(), Dtype(), std::move(attrs));
} else if (Dtype() == DType::kFloat32) {
txpr::AttributePayload attrs;
auto v = ToVector<float>();
attrs.tensor_data.assign(v.begin(), v.end());
return tape->AppendInstruction(txpr::OpCode::kConstant, {}, shape(), Dtype(), std::move(attrs));
} else if (Dtype() == DType::kFloat64) {
txpr::AttributePayload attrs;
attrs.tensor_data = ToVector<double>();
return tape->AppendInstruction(txpr::OpCode::kConstant, {}, shape(), Dtype(), std::move(attrs));
} else if (IsBoolean(Dtype())) {
txpr::AttributePayload attrs;
auto v = ToVector<bool>();
attrs.int_tensor_data.assign(v.begin(), v.end());
return tape->AppendInstruction(txpr::OpCode::kConstant, {}, shape(), Dtype(), std::move(attrs));
} else if (IsIntegral(Dtype())) {
txpr::AttributePayload attrs;
switch (Dtype()) {
case DType::kInt8: {
auto v = ToVector<int8_t>();
attrs.int_tensor_data.assign(v.begin(), v.end());
break;
}
case DType::kInt16: {
auto v = ToVector<int16_t>();
attrs.int_tensor_data.assign(v.begin(), v.end());
break;
}
case DType::kInt32: {
auto v = ToVector<int32_t>();
attrs.int_tensor_data.assign(v.begin(), v.end());
break;
}
case DType::kInt64: {
attrs.int_tensor_data = ToVector<int64_t>();
break;
}
case DType::kUInt8: {
auto v = ToVector<uint8_t>();
attrs.int_tensor_data.assign(v.begin(), v.end());
break;
}
case DType::kUInt16: {
auto v = ToVector<uint16_t>();
attrs.int_tensor_data.assign(v.begin(), v.end());
break;
}
case DType::kUInt32: {
auto v = ToVector<uint32_t>();
attrs.int_tensor_data.assign(v.begin(), v.end());
break;
}
case DType::kUInt64: {
auto v = ToVector<uint64_t>();
attrs.int_tensor_data.assign(v.begin(), v.end());
break;
}
case DType::kInt1: {
auto v = ToVector<tenzr::int1>();
for (auto x : v) {
attrs.int_tensor_data.push_back(static_cast<int64_t>(x));
}
break;
}
case DType::kUInt1: {
auto v = ToVector<tenzr::uint1>();
for (auto x : v) {
attrs.int_tensor_data.push_back(static_cast<int64_t>(x));
}
break;
}
case DType::kInt2: {
auto v = ToVector<tenzr::int2>();
for (auto x : v) {
attrs.int_tensor_data.push_back(static_cast<int64_t>(x));
}
break;
}
case DType::kUInt2: {
auto v = ToVector<tenzr::uint2>();
for (auto x : v) {
attrs.int_tensor_data.push_back(static_cast<int64_t>(x));
}
break;
}
case DType::kInt4: {
auto v = ToVector<tenzr::int4>();
for (auto x : v) {
attrs.int_tensor_data.push_back(static_cast<int64_t>(x));
}
break;
}
case DType::kUInt4: {
auto v = ToVector<tenzr::uint4>();
for (auto x : v) {
attrs.int_tensor_data.push_back(static_cast<int64_t>(x));
}
break;
}
default:
break;
}
return tape->AppendInstruction(txpr::OpCode::kConstant, {}, shape(), Dtype(), std::move(attrs));
} else if (IsFloatingPoint(Dtype())) {
txpr::AttributePayload attrs;
std::vector<double> out;
switch (Dtype()) {
case DType::kFloat16: {
auto v = ToVector<tenzr::float16>();
out.assign(v.begin(), v.end());
break;
}
case DType::kBFloat16: {
auto v = ToVector<tenzr::bfloat16>();
out.assign(v.begin(), v.end());
break;
}
case DType::kFloat8E3M4: {
auto v = ToVector<tenzr::float8_e3m4>();
out.assign(v.begin(), v.end());
break;
}
case DType::kFloat8E4M3: {
auto v = ToVector<tenzr::float8_e4m3>();
out.assign(v.begin(), v.end());
break;
}
case DType::kFloat8E4M3FN: {
auto v = ToVector<tenzr::float8_e4m3fn>();
out.assign(v.begin(), v.end());
break;
}
case DType::kFloat8E4M3FNUZ: {
auto v = ToVector<tenzr::float8_e4m3fnuz>();
out.assign(v.begin(), v.end());
break;
}
case DType::kFloat8E4M3B11FNUZ: {
auto v = ToVector<tenzr::float8_e4m3b11fnuz>();
out.assign(v.begin(), v.end());
break;
}
case DType::kFloat8E5M2: {
auto v = ToVector<tenzr::float8_e5m2>();
out.assign(v.begin(), v.end());
break;
}
case DType::kFloat8E5M2FNUZ: {
auto v = ToVector<tenzr::float8_e5m2fnuz>();
out.assign(v.begin(), v.end());
break;
}
case DType::kFloat8E8M0FNU: {
auto v = ToVector<tenzr::float8_e8m0fnu>();
out.assign(v.begin(), v.end());
break;
}
case DType::kFloat6E2M3FN: {
auto v = ToVector<tenzr::float6_e2m3fn>();
out.assign(v.begin(), v.end());
break;
}
case DType::kFloat6E3M2FN: {
auto v = ToVector<tenzr::float6_e3m2fn>();
out.assign(v.begin(), v.end());
break;
}
case DType::kFloat4E2M1FN: {
auto v = ToVector<tenzr::float4_e2m1fn>();
out.assign(v.begin(), v.end());
break;
}
default:
break;
}
attrs.tensor_data = std::move(out);
return tape->AppendInstruction(txpr::OpCode::kConstant, {}, shape(), Dtype(), std::move(attrs));
}
}
txpr::VarId vid = tape->AddVar(shape(), Dtype());
tape->MarkInput(vid);
return vid;
}
void Array::SyncHost() const {
TENZR_CHECK(node_ != nullptr, "Array is uninitialized");
auto var_node = std::dynamic_pointer_cast<VariableNode>(node_);
TENZR_CHECK(var_node != nullptr, "Cannot sync host on non-variable AST node");
var_node->SyncHost();
}
void Array::SyncDevice() const {
TENZR_CHECK(node_ != nullptr, "Array is uninitialized");
auto var_node = std::dynamic_pointer_cast<VariableNode>(node_);
TENZR_CHECK(var_node != nullptr, "Cannot sync device on non-variable AST node");
var_node->SyncDevice();
}
Array Array::Clone() const {
if (IsSymbolic() || txpr::IsGradTracingActive() || txpr::TracerContext::IsActive()) {
return *this;
}
TENZR_CHECK(node_ != nullptr, "Array is uninitialized");
auto var_node = std::dynamic_pointer_cast<VariableNode>(node_);
if (var_node) {
auto new_var = std::make_shared<VariableNode>(shape(), Dtype(), var_node->VariableName() + "_clone");
new_var->SetLayout(node_->Layout());
new_var->SetSharding(Sharding());
if (var_node->HasDeviceData() && !var_node->HasData()) {
var_node->SyncHost();
}
if (var_node->HasData()) {
auto src_buf = var_node->Data();
auto dst_buf = MemoryPool::Instance().Allocate(src_buf->SizeBytes());
std::memcpy(dst_buf->Data(), src_buf->Data(), src_buf->SizeBytes());
new_var->SetData(dst_buf);
}
return Array(new_var);
}
if (HasData()) {
auto new_var = std::make_shared<VariableNode>(shape(), Dtype(), "clone");
new_var->SetLayout(node_->Layout());
new_var->SetSharding(Sharding());
auto src_buf = Data();
auto dst_buf = MemoryPool::Instance().Allocate(src_buf->SizeBytes());
std::memcpy(dst_buf->Data(), src_buf->Data(), src_buf->SizeBytes());
new_var->SetData(dst_buf);
return Array(new_var);
}
return *this;
}
const StridedLayout& Array::Layout() const {
if (var_id_.IsValid()) {
static const StridedLayout default_layout;
return default_layout;
}
TENZR_CHECK(node_ != nullptr, "Array is uninitialized");
return node_->Layout();
}
bool Array::IsContiguous() const {
if (var_id_.IsValid()) {
return true;
}
return Layout().IsContiguous();
}
Array Array::Reshape(Shape new_shape) const {
if (txpr::TracerContext::IsActive() || IsSymbolic()) {
auto* tape = txpr::TracerContext::CurrentTape();
TENZR_CHECK(tape != nullptr, "Active tape required for symbolic operation");
txpr::VarId in_id = ObtainVarId(tape);
txpr::VarId out_v = tape->AppendInstruction(txpr::OpCode::kReshape, {in_id}, new_shape, Dtype());
return Array::Symbolic(out_v, std::move(new_shape), Dtype());
}
TENZR_CHECK(node_ != nullptr, "Array is uninitialized");
auto node = std::make_shared<ReshapeNode>(node_, std::move(new_shape));
return Array(std::move(node));
}
Array Array::Flatten() const { return Reshape(Shape({shape().NumElements()})); }
Array Array::Ravel() const { return Flatten(); }
Array Array::Transpose(int64_t dim0, int64_t dim1) const {
TENZR_CHECK(IsDefined(), "Array is uninitialized");
int64_t r = static_cast<int64_t>(shape().Rank());
TENZR_CHECK(r >= 2, "Transpose requires rank >= 2");
int64_t d0 = dim0 < 0 ? dim0 + r : dim0;
int64_t d1 = dim1 < 0 ? dim1 + r : dim1;
std::vector<int64_t> perm(static_cast<size_t>(r));
std::iota(perm.begin(), perm.end(), 0);
std::swap(perm[static_cast<size_t>(d0)], perm[static_cast<size_t>(d1)]);
return Permute(perm);
}
Array Array::Permute(const std::vector<int64_t>& axes) const {
if (txpr::TracerContext::IsActive() || IsSymbolic()) {
auto* tape = txpr::TracerContext::CurrentTape();
TENZR_CHECK(tape != nullptr, "Active tape required for symbolic operation");
txpr::VarId in_id = ObtainVarId(tape);
std::vector<int64_t> out_dims(shape().Rank());
for (size_t i = 0; i < shape().Rank(); ++i) {
out_dims[i] = shape()[static_cast<size_t>(axes[i])];
}
Shape out_shape(std::move(out_dims));
txpr::AttributePayload attrs;
attrs.permutation = axes;
txpr::VarId out_v = tape->AppendInstruction(txpr::OpCode::kTranspose, {in_id}, out_shape, Dtype(), std::move(attrs));
return Array::Symbolic(out_v, std::move(out_shape), Dtype());
}
TENZR_CHECK(node_ != nullptr, "Array is uninitialized");
auto var_node = std::dynamic_pointer_cast<VariableNode>(node_);
if (var_node && var_node->HasData() && !txpr::IsGradTracingActive()) {
auto new_lay = var_node->Layout().Permute(axes);
auto new_var = std::make_shared<VariableNode>(new_lay.shape(), Dtype(), var_node->VariableName() + "_permuted");
new_var->SetData(var_node->Data());
new_var->SetLayout(new_lay);
return Array(new_var);
}
return Array(std::make_shared<TransposeNode>(node_, axes));
}
Array Array::Squeeze(std::optional<int64_t> dim) const {
TENZR_CHECK(IsDefined(), "Array is uninitialized");
std::vector<int64_t> out_dims;
for (size_t i = 0; i < shape().Rank(); ++i) {
if (dim.has_value()) {
int64_t target_d = (*dim < 0) ? (*dim + static_cast<int64_t>(shape().Rank())) : *dim;
if (static_cast<int64_t>(i) == target_d) {
TENZR_CHECK(shape().Dim(i) == 1, "Cannot squeeze dimension with size != 1");
continue;
}
} else if (shape().Dim(i) == 1) {
continue;
}
out_dims.push_back(shape().Dim(i));
}
return Reshape(Shape(std::move(out_dims)));
}
Array Array::Unsqueeze(int64_t dim) const {
TENZR_CHECK(IsDefined(), "Array is uninitialized");
int64_t r = static_cast<int64_t>(shape().Rank());
int64_t target_d = (dim < 0) ? (dim + r + 1) : dim;
TENZR_CHECK(target_d >= 0 && target_d <= r, "unsqueeze dimension out of range");
std::vector<int64_t> out_dims;
for (int64_t i = 0; i < r; ++i) {
if (i == target_d) {
out_dims.push_back(1);
}
out_dims.push_back(shape().Dim(static_cast<size_t>(i)));
}
if (target_d == r) {
out_dims.push_back(1);
}
return Reshape(Shape(std::move(out_dims)));
}
Array Array::slice(const std::vector<Slice>& slices) const {
if (txpr::TracerContext::IsActive() || IsSymbolic()) {
auto* tape = txpr::TracerContext::CurrentTape();
TENZR_CHECK(tape != nullptr, "Active tape required for symbolic operation");
txpr::VarId in_id = ObtainVarId(tape);
size_t r = shape().Rank();
std::vector<int64_t> out_dims(r);
for (size_t i = 0; i < r; ++i) {
int64_t d = shape().Dim(i);
if (i < slices.size()) {
int64_t st = slices[i].start.value_or(0);
int64_t sp = slices[i].stop.value_or(d);
if (st < 0) {
st += d;
}
if (sp < 0) {
sp += d;
}
st = std::clamp(st, int64_t(0), d);
sp = std::clamp(sp, int64_t(0), d);
int64_t step = slices[i].step;
out_dims[i] = std::max(int64_t(0), (sp - st + step - 1) / step);
} else {
out_dims[i] = d;
}
}
Shape out_shape(std::move(out_dims));
txpr::AttributePayload attrs;
attrs.slices = slices;
txpr::VarId out_v = tape->AppendInstruction(txpr::OpCode::kSlice, {in_id}, out_shape, Dtype(), std::move(attrs));
return Array::Symbolic(out_v, std::move(out_shape), Dtype());
}
TENZR_CHECK(node_ != nullptr, "Array is uninitialized");
auto var_node = std::dynamic_pointer_cast<VariableNode>(node_);
if (var_node && var_node->HasData() && !txpr::IsGradTracingActive()) {
auto new_lay = var_node->Layout().slice(slices);
auto new_var = std::make_shared<VariableNode>(new_lay.shape(), Dtype(), var_node->VariableName() + "_sliced");
new_var->SetData(var_node->Data());
new_var->SetLayout(new_lay);
return Array(new_var);
}
return Array(std::make_shared<SliceNode>(node_, slices));
}
Array Array::BroadcastTo(const Shape& target_shape) const {
if (txpr::TracerContext::IsActive() || IsSymbolic()) {
auto* tape = txpr::TracerContext::CurrentTape();
TENZR_CHECK(tape != nullptr, "Active tape required for symbolic operation");
txpr::VarId in_id = ObtainVarId(tape);
size_t in_rank = shape().Rank();
size_t out_rank = target_shape.Rank();
std::vector<int64_t> dims(in_rank);
for (size_t i = 0; i < in_rank; ++i) {
dims[i] = static_cast<int64_t>((out_rank - in_rank) + i);
}
txpr::AttributePayload attrs;
attrs.axes = std::move(dims);
txpr::VarId out_v = tape->AppendInstruction(txpr::OpCode::kBroadcastTo, {in_id}, target_shape, Dtype(), std::move(attrs));
return Array::Symbolic(out_v, target_shape, Dtype());
}
TENZR_CHECK(node_ != nullptr, "Array is uninitialized");
auto var_node = std::dynamic_pointer_cast<VariableNode>(node_);
if (var_node && var_node->HasData() && !txpr::IsGradTracingActive()) {
auto new_lay = var_node->Layout().BroadcastTo(target_shape);
auto new_var = std::make_shared<VariableNode>(new_lay.shape(), Dtype(), var_node->VariableName() + "_broadcast");
new_var->SetData(var_node->Data());
new_var->SetLayout(new_lay);
return Array(new_var);
}
auto zero = Array::Zeros(target_shape, Dtype());
return *this + zero;
}
Array Array::AsStrided(const Shape& target_shape, const std::vector<int64_t>& strides, int64_t offset) const {
TENZR_CHECK(node_ != nullptr, "Array is uninitialized");
auto var_node = std::dynamic_pointer_cast<VariableNode>(node_);
TENZR_CHECK(var_node != nullptr && var_node->HasData(), "as_strided requires a tensor with bound host data");
StridedLayout new_lay(target_shape, strides, offset);
auto new_var = std::make_shared<VariableNode>(target_shape, Dtype(), var_node->VariableName() + "_strided");
new_var->SetData(var_node->Data());
new_var->SetLayout(new_lay);
return Array(new_var);
}
namespace {
template <typename T> Array MakeContiguousTyped(const Array& arr) {
auto vec = arr.ToVector<T>();
auto res = Array::Parameter(arr.shape(), arr.Dtype());
res.SetData(vec);
return res;
}
template <typename SrcT, typename DstT> void ConvertBufferElements(const SrcT* src, DstT* dst, size_t count) {
for (size_t i = 0; i < count; ++i) {
if constexpr (std::is_same_v<SrcT, bool>) {
double v = src[i] ? 1.0 : 0.0;
dst[i] = static_cast<DstT>(v);
} else if constexpr (std::is_same_v<DstT, bool>) {
double v = static_cast<double>(src[i]);
dst[i] = (v != 0.0);
} else {
double v = static_cast<double>(src[i]);
dst[i] = static_cast<DstT>(v);
}
}
}
template <typename SrcT> void DispatchDestType(const SrcT* src, const std::shared_ptr<PinnedBuffer>& dst_buf, DType dst_dtype, size_t count) {
switch (dst_dtype) {
case DType::kBool: {
uint8_t* dst = reinterpret_cast<uint8_t*>(dst_buf->Data());
for (size_t i = 0; i < count; ++i) {
if constexpr (std::is_same_v<SrcT, bool>) {
dst[i] = src[i] ? 1 : 0;
} else {
double v = static_cast<double>(src[i]);
dst[i] = (v != 0.0) ? 1 : 0;
}
}
break;
}
case DType::kInt8:
ConvertBufferElements(src, dst_buf->As<int8_t>(), count);
break;
case DType::kUInt8:
ConvertBufferElements(src, dst_buf->As<uint8_t>(), count);
break;
case DType::kInt16:
ConvertBufferElements(src, dst_buf->As<int16_t>(), count);
break;
case DType::kUInt16:
ConvertBufferElements(src, dst_buf->As<uint16_t>(), count);
break;
case DType::kInt32:
ConvertBufferElements(src, dst_buf->As<int32_t>(), count);
break;
case DType::kUInt32:
ConvertBufferElements(src, dst_buf->As<uint32_t>(), count);
break;
case DType::kInt64:
ConvertBufferElements(src, dst_buf->As<int64_t>(), count);
break;
case DType::kUInt64:
ConvertBufferElements(src, dst_buf->As<uint64_t>(), count);
break;
case DType::kFloat16:
ConvertBufferElements(src, dst_buf->As<tenzr::float16>(), count);
break;
case DType::kBFloat16:
ConvertBufferElements(src, dst_buf->As<tenzr::bfloat16>(), count);
break;
case DType::kFloat32:
ConvertBufferElements(src, dst_buf->As<float>(), count);
break;
case DType::kFloat64:
ConvertBufferElements(src, dst_buf->As<double>(), count);
break;
case DType::kFloat8E4M3FN:
ConvertBufferElements(src, dst_buf->As<tenzr::float8_e4m3fn>(), count);
break;
case DType::kFloat8E5M2:
ConvertBufferElements(src, dst_buf->As<tenzr::float8_e5m2>(), count);
break;
case DType::kFloat8E3M4:
ConvertBufferElements(src, dst_buf->As<tenzr::float8_e3m4>(), count);
break;
case DType::kFloat8E4M3:
ConvertBufferElements(src, dst_buf->As<tenzr::float8_e4m3>(), count);
break;
case DType::kFloat8E4M3FNUZ:
ConvertBufferElements(src, dst_buf->As<tenzr::float8_e4m3fnuz>(), count);
break;
case DType::kFloat8E4M3B11FNUZ:
ConvertBufferElements(src, dst_buf->As<tenzr::float8_e4m3b11fnuz>(), count);
break;
case DType::kFloat8E5M2FNUZ:
ConvertBufferElements(src, dst_buf->As<tenzr::float8_e5m2fnuz>(), count);
break;
case DType::kFloat8E8M0FNU:
ConvertBufferElements(src, dst_buf->As<tenzr::float8_e8m0fnu>(), count);
break;
case DType::kFloat6E2M3FN:
ConvertBufferElements(src, dst_buf->As<tenzr::float6_e2m3fn>(), count);
break;
case DType::kFloat6E3M2FN:
ConvertBufferElements(src, dst_buf->As<tenzr::float6_e3m2fn>(), count);
break;
case DType::kFloat4E2M1FN:
ConvertBufferElements(src, dst_buf->As<tenzr::float4_e2m1fn>(), count);
break;
case DType::kInt1:
ConvertBufferElements(src, dst_buf->As<tenzr::int1>(), count);
break;
case DType::kUInt1:
ConvertBufferElements(src, dst_buf->As<tenzr::uint1>(), count);
break;
case DType::kInt2:
ConvertBufferElements(src, dst_buf->As<tenzr::int2>(), count);
break;
case DType::kUInt2:
ConvertBufferElements(src, dst_buf->As<tenzr::uint2>(), count);
break;
case DType::kInt4:
ConvertBufferElements(src, dst_buf->As<tenzr::int4>(), count);
break;
case DType::kUInt4:
ConvertBufferElements(src, dst_buf->As<tenzr::uint4>(), count);
break;
default:
TENZR_THROW("Unsupported destination dtype for astype: " + std::string(ToString(dst_dtype)));
}
}
} // namespace
Array Array::Contiguous() const {
if (IsContiguous() && Layout().Offset() == 0) {
return *this;
}
switch (Dtype()) {
case DType::kBool:
return MakeContiguousTyped<bool>(*this);
case DType::kInt8:
return MakeContiguousTyped<int8_t>(*this);
case DType::kUInt8:
return MakeContiguousTyped<uint8_t>(*this);
case DType::kInt16:
return MakeContiguousTyped<int16_t>(*this);
case DType::kUInt16:
return MakeContiguousTyped<uint16_t>(*this);
case DType::kInt32:
return MakeContiguousTyped<int32_t>(*this);
case DType::kUInt32:
return MakeContiguousTyped<uint32_t>(*this);
case DType::kInt64:
return MakeContiguousTyped<int64_t>(*this);
case DType::kUInt64:
return MakeContiguousTyped<uint64_t>(*this);
case DType::kFloat16:
return MakeContiguousTyped<tenzr::float16>(*this);
case DType::kBFloat16:
return MakeContiguousTyped<tenzr::bfloat16>(*this);
case DType::kFloat32:
return MakeContiguousTyped<float>(*this);
case DType::kFloat64:
return MakeContiguousTyped<double>(*this);
case DType::kFloat8E4M3FN:
return MakeContiguousTyped<tenzr::float8_e4m3fn>(*this);
case DType::kFloat8E5M2:
return MakeContiguousTyped<tenzr::float8_e5m2>(*this);
case DType::kFloat8E3M4:
return MakeContiguousTyped<tenzr::float8_e3m4>(*this);
case DType::kFloat8E4M3:
return MakeContiguousTyped<tenzr::float8_e4m3>(*this);
case DType::kFloat8E4M3FNUZ:
return MakeContiguousTyped<tenzr::float8_e4m3fnuz>(*this);
case DType::kFloat8E4M3B11FNUZ:
return MakeContiguousTyped<tenzr::float8_e4m3b11fnuz>(*this);
case DType::kFloat8E5M2FNUZ:
return MakeContiguousTyped<tenzr::float8_e5m2fnuz>(*this);
case DType::kFloat8E8M0FNU:
return MakeContiguousTyped<tenzr::float8_e8m0fnu>(*this);
case DType::kFloat6E2M3FN:
return MakeContiguousTyped<tenzr::float6_e2m3fn>(*this);
case DType::kFloat6E3M2FN:
return MakeContiguousTyped<tenzr::float6_e3m2fn>(*this);
case DType::kFloat4E2M1FN:
return MakeContiguousTyped<tenzr::float4_e2m1fn>(*this);
case DType::kInt1:
return MakeContiguousTyped<tenzr::int1>(*this);
case DType::kUInt1:
return MakeContiguousTyped<tenzr::uint1>(*this);
case DType::kInt2:
return MakeContiguousTyped<tenzr::int2>(*this);
case DType::kUInt2:
return MakeContiguousTyped<tenzr::uint2>(*this);
case DType::kInt4:
return MakeContiguousTyped<tenzr::int4>(*this);
case DType::kUInt4:
return MakeContiguousTyped<tenzr::uint4>(*this);
default:
TENZR_THROW("Unsupported dtype in contiguous: " + std::string(tenzr::ToString(Dtype())));
}
}
Array Array::Astype(DType target_dtype) const {
if (Dtype() == target_dtype) {
return *this;
}
if (txpr::TracerContext::IsActive() || IsSymbolic()) {
auto* tape = txpr::TracerContext::CurrentTape();
TENZR_CHECK(tape != nullptr, "Active tape required for symbolic operation");
txpr::VarId in_id = ObtainVarId(tape);
txpr::AttributePayload attrs;
attrs.target_dtype = target_dtype;
txpr::VarId out_v = tape->AppendInstruction(txpr::OpCode::kConvert, {in_id}, shape(), target_dtype, std::move(attrs));
return Array::Symbolic(out_v, shape(), target_dtype);
}
TENZR_CHECK(node_ != nullptr, "Array is uninitialized");
if (HasData()) {
size_t count = static_cast<size_t>(shape().NumElements());
size_t bytes = count * SizeOf(target_dtype);
auto dst_buf = MemoryPool::Instance().Allocate(bytes);
switch (Dtype()) {
case DType::kBool: {
auto vec = ToVector<bool>();
std::vector<uint8_t> u8_vec(vec.begin(), vec.end());
DispatchDestType(u8_vec.data(), dst_buf, target_dtype, count);
break;
}
case DType::kInt8: {
auto vec = ToVector<int8_t>();
DispatchDestType(vec.data(), dst_buf, target_dtype, count);
break;
}
case DType::kUInt8: {
auto vec = ToVector<uint8_t>();
DispatchDestType(vec.data(), dst_buf, target_dtype, count);
break;
}
case DType::kInt16: {
auto vec = ToVector<int16_t>();
DispatchDestType(vec.data(), dst_buf, target_dtype, count);
break;
}
case DType::kUInt16: {
auto vec = ToVector<uint16_t>();
DispatchDestType(vec.data(), dst_buf, target_dtype, count);
break;
}
case DType::kInt32: {
auto vec = ToVector<int32_t>();
DispatchDestType(vec.data(), dst_buf, target_dtype, count);
break;
}
case DType::kUInt32: {
auto vec = ToVector<uint32_t>();
DispatchDestType(vec.data(), dst_buf, target_dtype, count);
break;
}
case DType::kInt64: {
auto vec = ToVector<int64_t>();
DispatchDestType(vec.data(), dst_buf, target_dtype, count);
break;
}
case DType::kUInt64: {
auto vec = ToVector<uint64_t>();
DispatchDestType(vec.data(), dst_buf, target_dtype, count);
break;
}
case DType::kFloat16: {
auto vec = ToVector<tenzr::float16>();
DispatchDestType(vec.data(), dst_buf, target_dtype, count);
break;
}
case DType::kBFloat16: {
auto vec = ToVector<tenzr::bfloat16>();
DispatchDestType(vec.data(), dst_buf, target_dtype, count);
break;
}
case DType::kFloat32: {
auto vec = ToVector<float>();
DispatchDestType(vec.data(), dst_buf, target_dtype, count);
break;
}
case DType::kFloat64: {
auto vec = ToVector<double>();
DispatchDestType(vec.data(), dst_buf, target_dtype, count);
break;
}
case DType::kFloat8E4M3FN: {
auto vec = ToVector<tenzr::float8_e4m3fn>();
DispatchDestType(vec.data(), dst_buf, target_dtype, count);
break;
}
case DType::kFloat8E5M2: {
auto vec = ToVector<tenzr::float8_e5m2>();
DispatchDestType(vec.data(), dst_buf, target_dtype, count);
break;
}
case DType::kFloat8E3M4: {
auto vec = ToVector<tenzr::float8_e3m4>();
DispatchDestType(vec.data(), dst_buf, target_dtype, count);
break;
}
case DType::kFloat8E4M3: {
auto vec = ToVector<tenzr::float8_e4m3>();
DispatchDestType(vec.data(), dst_buf, target_dtype, count);
break;
}
case DType::kFloat8E4M3FNUZ: {
auto vec = ToVector<tenzr::float8_e4m3fnuz>();
DispatchDestType(vec.data(), dst_buf, target_dtype, count);
break;
}
case DType::kFloat8E4M3B11FNUZ: {
auto vec = ToVector<tenzr::float8_e4m3b11fnuz>();
DispatchDestType(vec.data(), dst_buf, target_dtype, count);
break;
}
case DType::kFloat8E5M2FNUZ: {
auto vec = ToVector<tenzr::float8_e5m2fnuz>();
DispatchDestType(vec.data(), dst_buf, target_dtype, count);
break;
}
case DType::kFloat8E8M0FNU: {
auto vec = ToVector<tenzr::float8_e8m0fnu>();
DispatchDestType(vec.data(), dst_buf, target_dtype, count);
break;
}
case DType::kFloat6E2M3FN: {
auto vec = ToVector<tenzr::float6_e2m3fn>();
DispatchDestType(vec.data(), dst_buf, target_dtype, count);
break;
}
case DType::kFloat6E3M2FN: {
auto vec = ToVector<tenzr::float6_e3m2fn>();
DispatchDestType(vec.data(), dst_buf, target_dtype, count);
break;
}
case DType::kFloat4E2M1FN: {
auto vec = ToVector<tenzr::float4_e2m1fn>();
DispatchDestType(vec.data(), dst_buf, target_dtype, count);
break;
}
case DType::kInt1: {
auto vec = ToVector<tenzr::int1>();
DispatchDestType(vec.data(), dst_buf, target_dtype, count);
break;
}
case DType::kUInt1: {
auto vec = ToVector<tenzr::uint1>();
DispatchDestType(vec.data(), dst_buf, target_dtype, count);
break;
}
case DType::kInt2: {
auto vec = ToVector<tenzr::int2>();
DispatchDestType(vec.data(), dst_buf, target_dtype, count);
break;
}
case DType::kUInt2: {
auto vec = ToVector<tenzr::uint2>();
DispatchDestType(vec.data(), dst_buf, target_dtype, count);
break;
}
case DType::kInt4: {
auto vec = ToVector<tenzr::int4>();
DispatchDestType(vec.data(), dst_buf, target_dtype, count);
break;
}
case DType::kUInt4: {
auto vec = ToVector<tenzr::uint4>();
DispatchDestType(vec.data(), dst_buf, target_dtype, count);
break;
}
default:
TENZR_THROW("Unsupported source dtype in astype: " + std::string(tenzr::ToString(Dtype())));
}
auto res = Parameter(shape(), target_dtype);
res.BindData(dst_buf);
return res;
}
return Array(std::make_shared<ConvertNode>(node_, target_dtype));
}
Array Array::Matmul(const Array& other) const {
if (txpr::TracerContext::IsActive() || IsSymbolic() || other.IsSymbolic()) {
auto* tape = txpr::TracerContext::CurrentTape();
TENZR_CHECK(tape != nullptr, "Active tape required for symbolic operation");
Shape out_shape;
if (shape().Rank() == 2 && other.shape().Rank() == 2) {
TENZR_CHECK(shape()[1] == other.shape()[0], "Incompatible matmul dimensions");
out_shape = Shape({shape()[0], other.shape()[1]});
} else {
TENZR_CHECK(shape().Rank() == other.shape().Rank(), "Batched matmul requires equal ranks");
std::vector<int64_t> out_dims(shape().Rank());
for (size_t i = 0; i < shape().Rank() - 2; ++i) {
TENZR_CHECK(shape()[i] == other.shape()[i], "Batch dimension mismatch");
out_dims[i] = shape()[i];
}
TENZR_CHECK(shape()[shape().Rank() - 1] == other.shape()[other.shape().Rank() - 2], "Contracting dimension mismatch");
out_dims[shape().Rank() - 2] = shape()[shape().Rank() - 2];
out_dims[shape().Rank() - 1] = other.shape()[other.shape().Rank() - 1];
out_shape = Shape(std::move(out_dims));
}
DType common_dtype = PromoteTypes(Dtype(), other.Dtype());
txpr::VarId lhs_id = ObtainVarId(tape);
txpr::VarId rhs_id = other.ObtainVarId(tape);
txpr::VarId out_v = tape->AppendInstruction(txpr::OpCode::kMatmul, {lhs_id, rhs_id}, out_shape, common_dtype);
return Array::Symbolic(out_v, std::move(out_shape), common_dtype);
}
TENZR_CHECK(node_ != nullptr && other.node_ != nullptr, "Array is uninitialized");
auto node = std::make_shared<MatmulNode>(node_, other.node_);
return Array(std::move(node));
}
Array Array::Relu() const {
if (txpr::TracerContext::IsActive() || IsSymbolic()) {
auto* tape = txpr::TracerContext::CurrentTape();
TENZR_CHECK(tape != nullptr, "Active tape required for symbolic operation");
txpr::VarId in_id = ObtainVarId(tape);
txpr::VarId out_v = tape->AppendInstruction(txpr::OpCode::kRelu, {in_id}, shape(), Dtype());
return Array::Symbolic(out_v, shape(), Dtype());
}
TENZR_CHECK(node_ != nullptr, "Array is uninitialized");
auto node = std::make_shared<ReluNode>(node_);
return Array(std::move(node));
}
Array Array::sigmoid() const {
if (txpr::TracerContext::IsActive() || IsSymbolic()) {
auto* tape = txpr::TracerContext::CurrentTape();
TENZR_CHECK(tape != nullptr, "Active tape required for symbolic operation");
txpr::VarId in_id = ObtainVarId(tape);
txpr::VarId out_v = tape->AppendInstruction(txpr::OpCode::kSigmoid, {in_id}, shape(), Dtype());
return Array::Symbolic(out_v, shape(), Dtype());
}
TENZR_CHECK(node_ != nullptr, "Array is uninitialized");
auto node = std::make_shared<SigmoidNode>(node_);
return Array(std::move(node));
}
Array Array::CustomOp(std::string op_name, std::vector<Array> inputs) {
if (txpr::TracerContext::IsActive()) {
auto* tape = txpr::TracerContext::CurrentTape();
TENZR_CHECK(tape != nullptr, "Active tape required for custom_op tracing");
const auto* Definition = OpRegistry::Instance().Lookup(op_name);
TENZR_CHECK(Definition != nullptr, "Missing OpDef for custom op: " + op_name);
if (Definition->Category == OpCategory::kBinaryElementwise && inputs.size() >= 2) {
DType target_dtype = PromoteTypes(inputs[0].Dtype(), inputs[1].Dtype());
inputs[0] = (inputs[0].Dtype() == target_dtype) ? inputs[0] : inputs[0].Astype(target_dtype);
inputs[1] = (inputs[1].Dtype() == target_dtype) ? inputs[1] : inputs[1].Astype(target_dtype);
}
std::vector<txpr::VarId> in_ids;
in_ids.reserve(inputs.size());
std::vector<Shape> in_shapes;
std::vector<DType> in_dtypes;
in_shapes.reserve(inputs.size());
in_dtypes.reserve(inputs.size());
for (const auto& in : inputs) {
in_ids.push_back(in.ObtainVarId(tape));
in_shapes.push_back(in.shape());
in_dtypes.push_back(in.Dtype());
}
Shape out_shape = Definition->shape_fn ? Definition->shape_fn(in_shapes) : in_shapes[0];
DType out_dtype = Definition->dtype_fn ? Definition->dtype_fn(in_dtypes) : in_dtypes[0];
txpr::AttributePayload attrs;
attrs.custom_op_name = op_name;
txpr::VarId out_v = tape->AppendInstruction(txpr::OpCode::kCustomOp, std::move(in_ids), out_shape, out_dtype, std::move(attrs));
return Array::Symbolic(out_v, std::move(out_shape), out_dtype);
}
const auto* Definition = OpRegistry::Instance().Lookup(op_name);
if (Definition && Definition->Category == OpCategory::kBinaryElementwise && inputs.size() >= 2) {
DType target_dtype = PromoteTypes(inputs[0].Dtype(), inputs[1].Dtype());
inputs[0] = (inputs[0].Dtype() == target_dtype) ? inputs[0] : inputs[0].Astype(target_dtype);
inputs[1] = (inputs[1].Dtype() == target_dtype) ? inputs[1] : inputs[1].Astype(target_dtype);
}
std::vector<std::shared_ptr<Node>> nodes;
nodes.reserve(inputs.size());
for (const auto& t : inputs) {
TENZR_CHECK(t.IsDefined(), "Input tensor to custom_op is uninitialized");
nodes.push_back(t.node());
}
auto node = std::make_shared<CustomOpNode>(std::move(op_name), std::move(nodes), Definition);
return Array(std::move(node));
}
Array Array::CustomOp(std::string op_name) const {
if (txpr::TracerContext::IsActive() || IsSymbolic()) {
return CustomOp(std::move(op_name), {*this});
}
TENZR_CHECK(node_ != nullptr, "Array is uninitialized");
std::vector<std::shared_ptr<Node>> nodes = {node_};
auto node = std::make_shared<CustomOpNode>(std::move(op_name), std::move(nodes));
return Array(std::move(node));
}
Array Array::tanh() const { return CustomOp("tanh"); }
Array Array::Exp() const { return CustomOp("exp"); }
Array Array::Log() const { return CustomOp("log"); }
Array Array::Abs() const { return CustomOp("abs"); }
Array Array::Neg() const { return CustomOp("neg"); }
Array Array::Sin() const { return CustomOp("sin"); }
Array Array::Cos() const { return CustomOp("cos"); }
Array Array::Sqrt() const { return CustomOp("sqrt"); }
Array Array::Rsqrt() const { return CustomOp("rsqrt"); }
Array Array::Ceil() const { return CustomOp("ceil"); }
Array Array::Floor() const { return CustomOp("floor"); }
Array Array::Round() const { return CustomOp("round"); }
Array Array::Sign() const { return CustomOp("sign"); }
Array Array::Cbrt() const { return CustomOp("cbrt"); }
Array Array::Tan() const { return CustomOp("tan"); }
Array Array::Expm1() const { return CustomOp("expm1"); }
Array Array::Log1p() const { return CustomOp("log1p"); }
Array Array::IsFinite() const { return CustomOp("is_finite"); }
Array Array::LogicalNot() const { return CustomOp("logical_not"); }
Array Array::Silu() const {
TENZR_CHECK(node_ != nullptr, "Array is uninitialized");
return *this * sigmoid();
}
Array Array::Gelu() const {
TENZR_CHECK(node_ != nullptr, "Array is uninitialized");
auto x_sq = *this * *this;
auto x_cube = x_sq * *this;
auto c_coeff = Array::Full(shape(), 0.044715, Dtype());
auto c_sqrt_2_pi = Array::Full(shape(), 0.7978845608028654, Dtype());
auto c_one = Array::Full(shape(), 1.0, Dtype());
auto c_half = Array::Full(shape(), 0.5, Dtype());
auto inner = (*this + c_coeff * x_cube) * c_sqrt_2_pi;
return c_half * *this * (c_one + inner.tanh());
}
Array Array::Pow(const Array& other) const { return CustomOp("pow", {*this, other}); }
Array Array::Remainder(const Array& other) const { return CustomOp("remainder", {*this, other}); }
Array Array::Atan2(const Array& other) const { return CustomOp("atan2", {*this, other}); }
Array Array::Minimum(const Array& other) const { return CustomOp("minimum", {*this, other}); }
Array Array::Maximum(const Array& other) const { return CustomOp("maximum", {*this, other}); }
Array Array::BitwiseAnd(const Array& other) const { return CustomOp("bitwise_and", {*this, other}); }
Array Array::BitwiseOr(const Array& other) const { return CustomOp("bitwise_or", {*this, other}); }
Array Array::BitwiseXor(const Array& other) const { return CustomOp("bitwise_xor", {*this, other}); }
Array Array::ShiftLeft(const Array& other) const { return CustomOp("shift_left", {*this, other}); }
Array Array::ShiftRight(const Array& other) const { return CustomOp("shift_right_arithmetic", {*this, other}); }
Array Array::Clamp(const Array& min, const Array& max) const { return CustomOp("clamp", {min, *this, max}); }
Array Array::Equal(const Array& other) const { return CustomOp("compare_eq", {*this, other}); }
Array Array::NotEqual(const Array& other) const { return CustomOp("compare_ne", {*this, other}); }
Array Array::GreaterEqual(const Array& other) const { return CustomOp("compare_ge", {*this, other}); }
Array Array::GreaterThan(const Array& other) const { return CustomOp("compare_gt", {*this, other}); }
Array Array::LessEqual(const Array& other) const { return CustomOp("compare_le", {*this, other}); }
Array Array::LessThan(const Array& other) const { return CustomOp("compare_lt", {*this, other}); }
namespace {
Array TraceReduction(const Array& input, txpr::OpCode op, std::optional<std::vector<int64_t>> dims, bool keepdims) {
auto* tape = txpr::TracerContext::CurrentTape();
TENZR_CHECK(tape != nullptr, "Active tape required for symbolic reduction");
txpr::VarId in_id = input.ObtainVarId(tape);
std::vector<int64_t> d;
if (dims.has_value()) {
d = *dims;
} else {
d.resize(input.shape().Rank());
std::iota(d.begin(), d.end(), 0);
}
std::vector<int64_t> non_keep_dims;
std::unordered_map<int64_t, bool> is_red;
for (int64_t ax : d) {
if (ax < 0) {
ax += static_cast<int64_t>(input.shape().Rank());
}
is_red[ax] = true;
}
for (size_t i = 0; i < input.shape().Rank(); ++i) {
if (is_red[static_cast<int64_t>(i)]) {
if (keepdims) {
non_keep_dims.push_back(1);
}
} else {
non_keep_dims.push_back(input.shape()[i]);
}
}
Shape out_shape = (non_keep_dims.empty() && !keepdims) ? Shape({}) : Shape(std::move(non_keep_dims));
txpr::AttributePayload attrs;
attrs.axes = d;
attrs.keepdims = keepdims;
txpr::VarId out_v = tape->AppendInstruction(op, {in_id}, out_shape, input.Dtype(), std::move(attrs));
return Array::Symbolic(out_v, std::move(out_shape), input.Dtype());
}
} // namespace
Array Array::Sum(std::optional<std::vector<int64_t>> dims, bool keepdims) const {
if (Dtype() == DType::kBool) {
return Astype(DType::kInt32).Sum(dims, keepdims);
}
if (txpr::TracerContext::IsActive() || IsSymbolic()) {
return TraceReduction(*this, txpr::OpCode::kReduceSum, dims, keepdims);
}
TENZR_CHECK(node_ != nullptr, "Array is uninitialized");
std::vector<int64_t> d = dims.value_or(std::vector<int64_t>{});
return Array(std::make_shared<ReduceNode>(node_, ReduceOp::kSum, std::move(d), keepdims));
}
Array Array::Mean(std::optional<std::vector<int64_t>> dims, bool keepdims) const {
TENZR_CHECK(IsDefined(), "Array is uninitialized");
auto s = Sum(dims, keepdims);
int64_t count = 1;
if (dims.has_value() && !dims->empty()) {
int64_t r = static_cast<int64_t>(shape().Rank());
for (int64_t dim_val : *dims) {
int64_t c = (dim_val < 0) ? (dim_val + r) : dim_val;
TENZR_CHECK(c >= 0 && c < r, "Dimension out of range: " + std::to_string(dim_val));
count *= shape().Dim(static_cast<size_t>(c));
}
} else {
count = shape().NumElements();
}
auto count_arr = Array::Full(s.shape(), static_cast<double>(count), s.Dtype());
return s / count_arr;
}
Array Array::Max(std::optional<std::vector<int64_t>> dims, bool keepdims) const {
if (txpr::TracerContext::IsActive() || IsSymbolic()) {
return TraceReduction(*this, txpr::OpCode::kReduceMax, dims, keepdims);
}
TENZR_CHECK(node_ != nullptr, "Array is uninitialized");
std::vector<int64_t> d = dims.value_or(std::vector<int64_t>{});
return Array(std::make_shared<ReduceNode>(node_, ReduceOp::kMax, std::move(d), keepdims));
}
Array Array::Min(std::optional<std::vector<int64_t>> dims, bool keepdims) const {
if (txpr::TracerContext::IsActive() || IsSymbolic()) {
return TraceReduction(*this, txpr::OpCode::kReduceMin, dims, keepdims);
}
TENZR_CHECK(node_ != nullptr, "Array is uninitialized");
std::vector<int64_t> d = dims.value_or(std::vector<int64_t>{});
return Array(std::make_shared<ReduceNode>(node_, ReduceOp::kMin, std::move(d), keepdims));
}
Array Array::Prod(std::optional<std::vector<int64_t>> dims, bool keepdims) const {
if (txpr::TracerContext::IsActive() || IsSymbolic()) {
return TraceReduction(*this, txpr::OpCode::kReduceProd, dims, keepdims);
}
TENZR_CHECK(node_ != nullptr, "Array is uninitialized");
std::vector<int64_t> d = dims.value_or(std::vector<int64_t>{});
return Array(std::make_shared<ReduceNode>(node_, ReduceOp::kProd, std::move(d), keepdims));
}
Array Array::Var(std::optional<std::vector<int64_t>> dims, bool keepdims) const {
Array m = Mean(dims, true);
Array diff = *this - m;
return (diff * diff).Mean(dims, keepdims);
}
Array Array::StdDev(std::optional<std::vector<int64_t>> dims, bool keepdims) const { return Var(dims, keepdims).Sqrt(); }
Array Array::operator+(const Array& other) const {
if (txpr::TracerContext::IsActive() || IsSymbolic() || other.IsSymbolic()) {
auto* tape = txpr::TracerContext::CurrentTape();
TENZR_CHECK(tape != nullptr, "Active tape required for symbolic operation");
DType common_dtype = PromoteTypes(Dtype(), other.Dtype());
Shape out_shape = Shape::Broadcast(shape(), other.shape());
txpr::VarId lhs_id = ObtainVarId(tape);
txpr::VarId rhs_id = other.ObtainVarId(tape);
if (Dtype() != common_dtype) {
lhs_id = tape->AppendInstruction(txpr::OpCode::kConvert, {lhs_id}, shape(), common_dtype);
}
if (other.Dtype() != common_dtype) {
rhs_id = tape->AppendInstruction(txpr::OpCode::kConvert, {rhs_id}, other.shape(), common_dtype);
}
txpr::VarId out_v = tape->AppendInstruction(txpr::OpCode::kAdd, {lhs_id, rhs_id}, out_shape, common_dtype);
return Array::Symbolic(out_v, std::move(out_shape), common_dtype);
}
TENZR_CHECK(node_ != nullptr && other.node_ != nullptr, "Array is uninitialized");
DType common_dtype = PromoteTypes(Dtype(), other.Dtype());
Array lhs = (Dtype() == common_dtype) ? *this : this->Astype(common_dtype);
Array rhs = (other.Dtype() == common_dtype) ? other : other.Astype(common_dtype);
auto node = std::make_shared<AddNode>(lhs.node(), rhs.node());
return Array(std::move(node));
}
Array Array::operator*(const Array& other) const {
if (txpr::TracerContext::IsActive() || IsSymbolic() || other.IsSymbolic()) {
auto* tape = txpr::TracerContext::CurrentTape();
TENZR_CHECK(tape != nullptr, "Active tape required for symbolic operation");
DType common_dtype = PromoteTypes(Dtype(), other.Dtype());
Shape out_shape = Shape::Broadcast(shape(), other.shape());
txpr::VarId lhs_id = ObtainVarId(tape);
txpr::VarId rhs_id = other.ObtainVarId(tape);
if (Dtype() != common_dtype) {
lhs_id = tape->AppendInstruction(txpr::OpCode::kConvert, {lhs_id}, shape(), common_dtype);
}
if (other.Dtype() != common_dtype) {
rhs_id = tape->AppendInstruction(txpr::OpCode::kConvert, {rhs_id}, other.shape(), common_dtype);
}
txpr::VarId out_v = tape->AppendInstruction(txpr::OpCode::kMul, {lhs_id, rhs_id}, out_shape, common_dtype);
return Array::Symbolic(out_v, std::move(out_shape), common_dtype);
}
TENZR_CHECK(node_ != nullptr && other.node_ != nullptr, "Array is uninitialized");
DType common_dtype = PromoteTypes(Dtype(), other.Dtype());
Array lhs = (Dtype() == common_dtype) ? *this : this->Astype(common_dtype);
Array rhs = (other.Dtype() == common_dtype) ? other : other.Astype(common_dtype);
auto node = std::make_shared<MulNode>(lhs.node(), rhs.node());
return Array(std::move(node));
}
Array Array::operator-(const Array& other) const {
if (txpr::TracerContext::IsActive() || IsSymbolic() || other.IsSymbolic()) {
auto* tape = txpr::TracerContext::CurrentTape();
TENZR_CHECK(tape != nullptr, "Active tape required for symbolic operation");
DType common_dtype = PromoteTypes(Dtype(), other.Dtype());
Shape out_shape = Shape::Broadcast(shape(), other.shape());
txpr::VarId lhs_id = ObtainVarId(tape);
txpr::VarId rhs_id = other.ObtainVarId(tape);
if (Dtype() != common_dtype) {
lhs_id = tape->AppendInstruction(txpr::OpCode::kConvert, {lhs_id}, shape(), common_dtype);
}
if (other.Dtype() != common_dtype) {
rhs_id = tape->AppendInstruction(txpr::OpCode::kConvert, {rhs_id}, other.shape(), common_dtype);
}
txpr::VarId out_v = tape->AppendInstruction(txpr::OpCode::kSub, {lhs_id, rhs_id}, out_shape, common_dtype);
return Array::Symbolic(out_v, std::move(out_shape), common_dtype);
}
TENZR_CHECK(node_ != nullptr && other.node_ != nullptr, "Array is uninitialized");
DType common_dtype = PromoteTypes(Dtype(), other.Dtype());
Array lhs = (Dtype() == common_dtype) ? *this : this->Astype(common_dtype);
Array rhs = (other.Dtype() == common_dtype) ? other : other.Astype(common_dtype);
auto node = std::make_shared<SubNode>(lhs.node(), rhs.node());
return Array(std::move(node));
}
Array Array::operator/(const Array& other) const {
if (txpr::TracerContext::IsActive() || IsSymbolic() || other.IsSymbolic()) {
auto* tape = txpr::TracerContext::CurrentTape();
TENZR_CHECK(tape != nullptr, "Active tape required for symbolic operation");
DType common_dtype = PromoteTypes(Dtype(), other.Dtype());
Shape out_shape = Shape::Broadcast(shape(), other.shape());
txpr::VarId lhs_id = ObtainVarId(tape);
txpr::VarId rhs_id = other.ObtainVarId(tape);
if (Dtype() != common_dtype) {
lhs_id = tape->AppendInstruction(txpr::OpCode::kConvert, {lhs_id}, shape(), common_dtype);
}
if (other.Dtype() != common_dtype) {
rhs_id = tape->AppendInstruction(txpr::OpCode::kConvert, {rhs_id}, other.shape(), common_dtype);
}
txpr::VarId out_v = tape->AppendInstruction(txpr::OpCode::kDiv, {lhs_id, rhs_id}, out_shape, common_dtype);
return Array::Symbolic(out_v, std::move(out_shape), common_dtype);
}
TENZR_CHECK(node_ != nullptr && other.node_ != nullptr, "Array is uninitialized");
DType common_dtype = PromoteTypes(Dtype(), other.Dtype());
Array lhs = (Dtype() == common_dtype) ? *this : this->Astype(common_dtype);
Array rhs = (other.Dtype() == common_dtype) ? other : other.Astype(common_dtype);
auto node = std::make_shared<DivNode>(lhs.node(), rhs.node());
return Array(std::move(node));
}
Array Array::operator==(const Array& other) const { return Equal(other); }
Array Array::operator!=(const Array& other) const { return NotEqual(other); }
Array Array::operator<(const Array& other) const { return LessThan(other); }
Array Array::operator<=(const Array& other) const { return LessEqual(other); }
Array Array::operator>(const Array& other) const { return GreaterThan(other); }
Array Array::operator>=(const Array& other) const { return GreaterEqual(other); }
std::string Array::ToString() const {
if (!node_) {
return "Array(empty)";
}
std::ostringstream oss;
oss << "Array(" << node_->OpName() << ", shape=" << shape().ToString() << ", dtype=" << tenzr::ToString(Dtype()) << ")";
return oss.str();
}
namespace {
ExecutionEngine& DefaultEvalEngine() {
static thread_local ExecutionEngine engine;
return engine;
}
} // namespace
Array Array::Eval() const {
TENZR_CHECK(!IsSymbolic(), "Cannot evaluate a purely symbolic Array outside of a compiled JIT context");
if (HasData()) {
return *this;
}
TENZR_CHECK(IsDefined(), "Cannot evaluate undefined tensor");
ExecutionResult res = DefaultEvalEngine().Execute(*this);
return res.AsArray();
}
Array Array::Materialize() const {
TENZR_CHECK(IsDefined(), "Cannot materialize an undefined Array");
TENZR_CHECK(!IsSymbolic(), "Cannot materialize a purely symbolic Array outside of a compiled JIT context");
if (HasData()) {
return *this;
}
return Eval();
}
namespace {
template <typename T> std::vector<std::string> FormatTyped(const Array& t) {
auto v = t.ToVector<T>();
std::vector<std::string> res;
res.reserve(v.size());
for (const auto& x : v) {
std::ostringstream ss;
if constexpr (std::is_same_v<T, int8_t>) {
ss << static_cast<int>(x);
} else if constexpr (std::is_same_v<T, uint8_t>) {
ss << static_cast<unsigned int>(x);
} else if constexpr (std::is_same_v<T, tenzr::float16> || std::is_same_v<T, tenzr::bfloat16> || std::is_same_v<T, tenzr::float8_e4m3fn> ||
std::is_same_v<T, tenzr::float8_e5m2> || std::is_same_v<T, tenzr::float8_e3m4> || std::is_same_v<T, tenzr::float8_e4m3> ||
std::is_same_v<T, tenzr::float8_e4m3fnuz> || std::is_same_v<T, tenzr::float8_e4m3b11fnuz> ||
std::is_same_v<T, tenzr::float8_e5m2fnuz> || std::is_same_v<T, tenzr::float8_e8m0fnu> || std::is_same_v<T, tenzr::float6_e2m3fn> ||
std::is_same_v<T, tenzr::float6_e3m2fn> || std::is_same_v<T, tenzr::float4_e2m1fn>) {
ss << static_cast<float>(x);
} else {
ss << x;
}
res.push_back(ss.str());
}
return res;
}
std::vector<std::string> FormatTensorElements(const Array& t) {
switch (t.Dtype()) {
case DType::kBool: {
auto v = t.ToVector<bool>();
std::vector<std::string> res;
res.reserve(v.size());
for (bool b : v) {
res.push_back(b ? "true" : "false");
}
return res;
}
case DType::kInt8:
return FormatTyped<int8_t>(t);
case DType::kUInt8:
return FormatTyped<uint8_t>(t);
case DType::kInt16:
return FormatTyped<int16_t>(t);
case DType::kUInt16:
return FormatTyped<uint16_t>(t);
case DType::kInt32:
return FormatTyped<int32_t>(t);
case DType::kUInt32:
return FormatTyped<uint32_t>(t);
case DType::kInt64:
return FormatTyped<int64_t>(t);
case DType::kUInt64:
return FormatTyped<uint64_t>(t);
case DType::kFloat16:
return FormatTyped<tenzr::float16>(t);
case DType::kBFloat16:
return FormatTyped<tenzr::bfloat16>(t);
case DType::kFloat32:
return FormatTyped<float>(t);
case DType::kFloat64:
return FormatTyped<double>(t);
case DType::kFloat8E4M3FN:
return FormatTyped<tenzr::float8_e4m3fn>(t);
case DType::kFloat8E5M2:
return FormatTyped<tenzr::float8_e5m2>(t);
case DType::kFloat8E3M4:
return FormatTyped<tenzr::float8_e3m4>(t);
case DType::kFloat8E4M3:
return FormatTyped<tenzr::float8_e4m3>(t);
case DType::kFloat8E4M3FNUZ:
return FormatTyped<tenzr::float8_e4m3fnuz>(t);
case DType::kFloat8E4M3B11FNUZ:
return FormatTyped<tenzr::float8_e4m3b11fnuz>(t);
case DType::kFloat8E5M2FNUZ:
return FormatTyped<tenzr::float8_e5m2fnuz>(t);
case DType::kFloat8E8M0FNU:
return FormatTyped<tenzr::float8_e8m0fnu>(t);
case DType::kFloat6E2M3FN:
return FormatTyped<tenzr::float6_e2m3fn>(t);
case DType::kFloat6E3M2FN:
return FormatTyped<tenzr::float6_e3m2fn>(t);
case DType::kFloat4E2M1FN:
return FormatTyped<tenzr::float4_e2m1fn>(t);
case DType::kInt1:
return FormatTyped<tenzr::int1>(t);
case DType::kUInt1:
return FormatTyped<tenzr::uint1>(t);
case DType::kInt2:
return FormatTyped<tenzr::int2>(t);
case DType::kUInt2:
return FormatTyped<tenzr::uint2>(t);
case DType::kInt4:
return FormatTyped<tenzr::int4>(t);
case DType::kUInt4:
return FormatTyped<tenzr::uint4>(t);
default:
TENZR_THROW("Unsupported dtype in FormatTensorElements: " + std::string(ToString(t.Dtype())));
}
}
} // namespace
std::ostream& operator<<(std::ostream& os, const Array& t) {
if (!t.IsDefined()) {
return os << "Array(undefined)";
}
if (t.IsSymbolic()) {
return os << "Array(symbolic=v" << t.SymbolicVarId().id << ", shape=" << t.shape().ToString() << ", dtype=" << tenzr::ToString(t.Dtype()) << ")";
}
Array evaluated = t;
if (!evaluated.HasData()) {
bool evaluated_ok = false;
#if TENZR_EXCEPTIONS_ENABLED
try {
evaluated = t.Eval();
evaluated_ok = true;
} catch (...) {
evaluated_ok = false;
}
#else
evaluated_ok = evaluated.HasData();
#endif
if (!evaluated_ok) {
return os << "Array(lazy_op=" << t.node()->OpName() << ", shape=" << t.shape().ToString() << ", dtype=" << tenzr::ToString(t.Dtype()) << ")";
}
}
os << "Array(";
const auto& s = evaluated.shape();
if (s.NumElements() == 0) {
os << "[]";
} else if (s.Rank() == 1) {
auto vec = FormatTensorElements(evaluated);
os << "[";
size_t n = vec.size();
size_t max_print = 8;
if (n <= max_print) {
for (size_t i = 0; i < n; ++i) {
os << (i > 0 ? ", " : "") << vec[i];
}
} else {
for (size_t i = 0; i < 3; ++i) {
os << (i > 0 ? ", " : "") << vec[i];
}
os << ", ...";
for (size_t i = n - 3; i < n; ++i) {
os << ", " << vec[i];
}
}
os << "]";
} else if (s.Rank() == 2) {
auto vec = FormatTensorElements(evaluated);
int64_t rows = s.Dim(0);
int64_t cols = s.Dim(1);
os << "[\n";
int64_t max_rows = 6;
for (int64_t r = 0; r < rows; ++r) {
if (rows > max_rows && r == 3) {
os << " ...,\n";
r = rows - 3;
}
os << " [";
int64_t max_cols = 6;
for (int64_t c = 0; c < cols; ++c) {
if (cols > max_cols && c == 3) {
os << ", ...";
c = cols - 3;
}
os << (c > 0 ? ", " : "") << vec[static_cast<size_t>(r * cols + c)];
}
os << "]" << (r + 1 < rows ? ",\n" : "\n");
}
os << "]";
} else {
auto vec = FormatTensorElements(evaluated);
os << "[";
size_t n = vec.size();
size_t print_cnt = std::min(n, size_t(6));
for (size_t i = 0; i < print_cnt; ++i) {
os << (i > 0 ? ", " : "") << vec[i];
}
if (n > print_cnt) {
os << ", ...";
}
os << "]";
}
os << ", shape=" << s.ToString() << ", dtype=" << tenzr::ToString(evaluated.Dtype()) << ")";
return os;
}
} // namespace tenzr