// 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/node.h"
#include <chrono>
#include <cstring>
#include <numeric>
#include "tenzr/common/error.h"
#include "tenzr/common/logging.h"
#include "tenzr/common/trace.h"
#include "tenzr/frontend/array.h"
#include "tenzr/frontend/dtype.h"
#include "tenzr/frontend/op_registry.h"
#include "tenzr/runtime/device_buffer.h"
#include "tenzr/runtime/hardware_manager.h"
#include "tenzr/runtime/pjrt_error.h"
namespace tenzr {
uint64_t Node::NextId() {
static std::atomic<uint64_t> counter{1};
return counter.fetch_add(1, std::memory_order_relaxed);
}
Node::Node(NodeKind kind, Shape shape, DType dtype, ShardingSpec sharding)
: kind_(kind), shape_(std::move(shape)), dtype_(dtype), sharding_(std::move(sharding)), layout_(shape_), id_(NextId()) {}
// VariableNode
VariableNode::VariableNode(Shape shape, DType dtype, std::string name, ShardingSpec sharding)
: Node(NodeKind::kVariable, std::move(shape), dtype, std::move(sharding)), name_(std::move(name)) {
if (name_.empty()) {
name_ = "var_" + std::to_string(id_);
}
}
void VariableNode::SetData(std::shared_ptr<PinnedBuffer> data) {
data_ = std::move(data);
device_data_ = nullptr;
}
const std::shared_ptr<PinnedBuffer>& VariableNode::Data() const {
if (!data_ && device_data_) {
const_cast<VariableNode*>(this)->SyncHost();
}
return data_;
}
bool VariableNode::HasData() const { return data_ != nullptr || device_data_ != nullptr; }
void VariableNode::SetDeviceData(std::shared_ptr<DeviceBuffer> device_data) { device_data_ = std::move(device_data); }
const std::shared_ptr<DeviceBuffer>& VariableNode::DeviceData() const { return device_data_; }
bool VariableNode::HasDeviceData() const { return device_data_ != nullptr; }
void VariableNode::SyncHost() {
if (data_) {
return;
}
TENZR_TRACE_SCOPE_CAT("transfer", "PJRT::D2H_Transfer");
TENZR_CHECK(device_data_ != nullptr, "Cannot sync host: variable has no device data");
auto& hw = HardwareManager::Instance();
const PJRT_Api* api = hw.Api();
TENZR_CHECK(api != nullptr, "HardwareManager API uninitialized");
size_t out_num_elems = static_cast<size_t>(shape_.NumElements());
size_t out_bytes = out_num_elems * SizeOf(dtype_);
auto host_buf = MemoryPool::Instance().Allocate(out_bytes);
PJRT_Buffer_ToHostBuffer_Args thb_args;
std::memset(&thb_args, 0, sizeof(thb_args));
thb_args.struct_size = PJRT_Buffer_ToHostBuffer_Args_STRUCT_SIZE;
thb_args.src = device_data_->Handle();
thb_args.host_layout = nullptr;
thb_args.dst = host_buf->Data();
thb_args.dst_size = out_bytes;
const auto transfer_start = std::chrono::steady_clock::now();
PJRT_Error* err = api->PJRT_Buffer_ToHostBuffer(&thb_args);
if (err) {
std::string err_msg = PjrtErrorMessage(api, err);
PjrtErrorDestroy(api, err);
TENZR_THROW("sync_host failed: " + err_msg);
}
if (thb_args.event) {
PJRT_Event_Await_Args await_args{PJRT_Event_Await_Args_STRUCT_SIZE, nullptr, thb_args.event};
PJRT_Error* await_err = api->PJRT_Event_Await(&await_args);
if (await_err != nullptr) {
std::string err_msg = PjrtErrorMessage(api, await_err);
PjrtErrorDestroy(api, await_err);
TENZR_LOG_WARN << "PJRT_Event_Await failed: " << err_msg << "\n";
}
PJRT_Event_Destroy_Args d_args{PJRT_Event_Destroy_Args_STRUCT_SIZE, nullptr, thb_args.event};
(void)api->PJRT_Event_Destroy(&d_args);
}
const size_t device_bytes = device_data_->DeviceSizeBytes();
const size_t moved_bytes = device_bytes > 0 ? device_bytes : out_bytes;
const auto transfer_us = std::chrono::duration_cast<std::chrono::microseconds>(std::chrono::steady_clock::now() - transfer_start).count();
profiler::RecordMemoryTransfer("PJRT::D2H", moved_bytes, false, static_cast<int64_t>(transfer_us));
data_ = std::move(host_buf);
}
void VariableNode::SyncDevice() {
TENZR_TRACE_SCOPE_CAT("transfer", "PJRT::H2D_Transfer");
if (device_data_) {
return;
}
TENZR_CHECK(data_ != nullptr, "Cannot sync device: variable has no host data");
if (!layout_.IsContiguous() || layout_.Offset() != 0) {
Array self_arr(shared_from_this());
Array contig = self_arr.Contiguous();
auto contig_node = std::dynamic_pointer_cast<VariableNode>(contig.node());
if (contig_node && contig_node->HasData()) {
data_ = contig_node->Data();
layout_ = contig_node->Layout();
}
}
auto& hw = HardwareManager::Instance();
if (!hw.IsInitialized()) {
hw.Initialize();
}
const PJRT_Api* api = hw.Api();
PJRT_Client* client = hw.Client();
PJRT_Device* device = hw.DefaultDevice();
TENZR_CHECK(api != nullptr && client != nullptr, "HardwareManager client uninitialized");
if (dtype_ == DType::kInt1 || dtype_ == DType::kUInt1) {
TENZR_THROW(
"1-bit dtypes (int1/uint1) cannot be transferred to PJRT: the plugin's S1/U1 buffer "
"layout is unsupported; cast to int8 (astype) before device execution");
}
PJRT_Client_BufferFromHostBuffer_Args b_args;
std::memset(&b_args, 0, sizeof(b_args));
b_args.struct_size = PJRT_Client_BufferFromHostBuffer_Args_STRUCT_SIZE;
b_args.client = client;
b_args.data = data_->Data();
b_args.type = ToPjrtBufferType(dtype_);
b_args.dims = shape_.Dims().data();
b_args.num_dims = shape_.Rank();
b_args.host_buffer_semantics = PJRT_HostBufferSemantics_kImmutableUntilTransferCompletes;
b_args.device = device;
const auto transfer_start = std::chrono::steady_clock::now();
PJRT_Error* err = api->PJRT_Client_BufferFromHostBuffer(&b_args);
if (err) {
std::string err_msg = PjrtErrorMessage(api, err);
PjrtErrorDestroy(api, err);
TENZR_THROW("sync_device failed: " + err_msg);
}
if (b_args.done_with_host_buffer) {
PJRT_Event_Await_Args await_args{PJRT_Event_Await_Args_STRUCT_SIZE, nullptr, b_args.done_with_host_buffer};
PJRT_Error* await_err = api->PJRT_Event_Await(&await_args);
if (await_err != nullptr) {
std::string err_msg = PjrtErrorMessage(api, await_err);
PjrtErrorDestroy(api, await_err);
TENZR_LOG_WARN << "PJRT_Event_Await failed: " << err_msg << "\n";
}
PJRT_Event_Destroy_Args d_args{PJRT_Event_Destroy_Args_STRUCT_SIZE, nullptr, b_args.done_with_host_buffer};
(void)api->PJRT_Event_Destroy(&d_args);
}
device_data_ = std::make_shared<DeviceBuffer>(b_args.buffer, api);
const size_t moved_bytes = static_cast<size_t>(shape_.NumElements()) * SizeOf(dtype_);
const auto transfer_us = std::chrono::duration_cast<std::chrono::microseconds>(std::chrono::steady_clock::now() - transfer_start).count();
profiler::RecordMemoryTransfer("PJRT::H2D", moved_bytes, true, static_cast<int64_t>(transfer_us));
}
// AddNode
AddNode::AddNode(std::shared_ptr<Node> lhs, std::shared_ptr<Node> rhs) : Node(NodeKind::kAdd, Shape::Broadcast(lhs->shape(), rhs->shape()), lhs->Dtype()) {
TENZR_CHECK(lhs != nullptr && rhs != nullptr, "AddNode inputs cannot be null");
TENZR_CHECK(lhs->Dtype() == rhs->Dtype(),
"AddNode inputs must have identical dtype: " + std::string(ToString(lhs->Dtype())) + " vs " + std::string(ToString(rhs->Dtype())));
inputs_ = {std::move(lhs), std::move(rhs)};
}
// MulNode
MulNode::MulNode(std::shared_ptr<Node> lhs, std::shared_ptr<Node> rhs) : Node(NodeKind::kMul, Shape::Broadcast(lhs->shape(), rhs->shape()), lhs->Dtype()) {
TENZR_CHECK(lhs != nullptr && rhs != nullptr, "MulNode inputs cannot be null");
TENZR_CHECK(lhs->Dtype() == rhs->Dtype(),
"MulNode inputs must have identical dtype: " + std::string(ToString(lhs->Dtype())) + " vs " + std::string(ToString(rhs->Dtype())));
inputs_ = {std::move(lhs), std::move(rhs)};
}
// SubNode
SubNode::SubNode(std::shared_ptr<Node> lhs, std::shared_ptr<Node> rhs) : Node(NodeKind::kSub, Shape::Broadcast(lhs->shape(), rhs->shape()), lhs->Dtype()) {
TENZR_CHECK(lhs != nullptr && rhs != nullptr, "SubNode inputs cannot be null");
TENZR_CHECK(lhs->Dtype() == rhs->Dtype(),
"SubNode inputs must have identical dtype: " + std::string(ToString(lhs->Dtype())) + " vs " + std::string(ToString(rhs->Dtype())));
inputs_ = {std::move(lhs), std::move(rhs)};
}
// DivNode
DivNode::DivNode(std::shared_ptr<Node> lhs, std::shared_ptr<Node> rhs) : Node(NodeKind::kDiv, Shape::Broadcast(lhs->shape(), rhs->shape()), lhs->Dtype()) {
TENZR_CHECK(lhs != nullptr && rhs != nullptr, "DivNode inputs cannot be null");
TENZR_CHECK(lhs->Dtype() == rhs->Dtype(),
"DivNode inputs must have identical dtype: " + std::string(ToString(lhs->Dtype())) + " vs " + std::string(ToString(rhs->Dtype())));
inputs_ = {std::move(lhs), std::move(rhs)};
}
// MatmulNode
MatmulNode::MatmulNode(std::shared_ptr<Node> lhs, std::shared_ptr<Node> rhs) : Node(NodeKind::kMatmul, Shape{}, lhs->Dtype()) {
TENZR_CHECK(lhs != nullptr && rhs != nullptr, "MatmulNode inputs cannot be null");
TENZR_CHECK(lhs->Dtype() == rhs->Dtype(),
"MatmulNode inputs must have identical dtype: " + std::string(ToString(lhs->Dtype())) + " vs " + std::string(ToString(rhs->Dtype())));
TENZR_CHECK(lhs->shape().Rank() >= 2 && rhs->shape().Rank() >= 2,
"Matmul requires rank >= 2, got lhs=" + std::to_string(lhs->shape().Rank()) + ", rhs=" + std::to_string(rhs->shape().Rank()));
int64_t lhs_k = lhs->shape().Dim(lhs->shape().Rank() - 1);
int64_t rhs_k = rhs->shape().Dim(rhs->shape().Rank() - 2);
TENZR_CHECK(lhs_k == rhs_k, "Incompatible contracting dimensions for Matmul: " + std::to_string(lhs_k) + " vs " + std::to_string(rhs_k));
int64_t m = lhs->shape().Dim(lhs->shape().Rank() - 2);
int64_t n = rhs->shape().Dim(rhs->shape().Rank() - 1);
std::vector<int64_t> out_dims;
// For 2D: [M, N]
if (lhs->shape().Rank() == 2 && rhs->shape().Rank() == 2) {
out_dims = {m, n};
} else {
// Handle batch dimensions if higher rank
for (size_t i = 0; i < lhs->shape().Rank() - 2; ++i) {
out_dims.push_back(lhs->shape().Dim(i));
}
out_dims.push_back(m);
out_dims.push_back(n);
}
shape_ = Shape(out_dims);
inputs_ = {std::move(lhs), std::move(rhs)};
}
ConvolutionNode::ConvolutionNode(std::shared_ptr<Node> input, std::shared_ptr<Node> kernel, std::vector<int64_t> window_strides,
std::vector<int64_t> lhs_dilation, std::vector<int64_t> rhs_dilation, std::vector<int64_t> padding_low,
std::vector<int64_t> padding_high)
: Node(NodeKind::kConvolution, Shape{}, input ? input->Dtype() : DType::kFloat32, input ? input->Sharding() : ShardingSpec::Replicated()),
window_strides_(std::move(window_strides)),
lhs_dilation_(std::move(lhs_dilation)),
rhs_dilation_(std::move(rhs_dilation)),
padding_low_(std::move(padding_low)),
padding_high_(std::move(padding_high)) {
TENZR_CHECK(input != nullptr && kernel != nullptr, "ConvolutionNode inputs cannot be null");
inputs_ = {std::move(input), std::move(kernel)};
const Shape& in = inputs_[0]->shape();
const Shape& kern = inputs_[1]->shape();
TENZR_CHECK(in.Rank() == 4 && kern.Rank() == 4, "ConvolutionNode expects NCHW input and OIHW kernel");
std::vector<int64_t> out_dims = {in.Dim(0), kern.Dim(0), 0, 0};
for (size_t d = 0; d < 2; ++d) {
const int64_t in_d = in.Dim(2 + d);
const int64_t k_d = kern.Dim(2 + d);
const int64_t lhs_dil = lhs_dilation_.empty() ? 1 : lhs_dilation_[d];
const int64_t rhs_dil = rhs_dilation_.empty() ? 1 : rhs_dilation_[d];
const int64_t stride = window_strides_.empty() ? 1 : window_strides_[d];
const int64_t lo = padding_low_.empty() ? 0 : padding_low_[d];
const int64_t hi = padding_high_.empty() ? 0 : padding_high_[d];
const int64_t lhs_dilated = (in_d - 1) * lhs_dil + 1;
const int64_t rhs_dilated = (k_d - 1) * rhs_dil + 1;
const int64_t numer = lhs_dilated + lo + hi - rhs_dilated;
TENZR_CHECK(numer >= 0, "Convolution kernel is larger than the padded input");
out_dims[2 + d] = numer / stride + 1;
}
shape_ = Shape(out_dims);
}
// ReshapeNode
ReshapeNode::ReshapeNode(std::shared_ptr<Node> input, Shape new_shape) : Node(NodeKind::kReshape, new_shape, input->Dtype()) {
TENZR_CHECK(input != nullptr, "Reshape input cannot be null");
TENZR_CHECK(input->shape().NumElements() == new_shape.NumElements(),
"Reshape elements must match: " + std::to_string(input->shape().NumElements()) + " vs " + std::to_string(new_shape.NumElements()));
inputs_ = {std::move(input)};
}
// TransposeNode
TransposeNode::TransposeNode(std::shared_ptr<Node> input, std::vector<int64_t> permutation)
: Node(NodeKind::kTranspose, Shape({}), input->Dtype(), input->Sharding()), permutation_(std::move(permutation)) {
TENZR_CHECK(input != nullptr, "Transpose input cannot be null");
const size_t r = input->shape().Rank();
TENZR_CHECK(permutation_.size() == r, "Permutation size must match rank");
std::vector<int64_t> new_dims(r);
for (size_t i = 0; i < r; ++i) {
int64_t axis = permutation_[i];
TENZR_CHECK(axis >= 0 && axis < static_cast<int64_t>(r), "Permutation axis out of range");
new_dims[i] = input->shape().Dim(static_cast<size_t>(axis));
}
shape_ = Shape(new_dims);
inputs_ = {std::move(input)};
}
// SliceNode
SliceNode::SliceNode(std::shared_ptr<Node> input, std::vector<Slice> slices)
: Node(NodeKind::kSlice, Shape({}), input->Dtype(), input->Sharding()), slices_(std::move(slices)) {
TENZR_CHECK(input != nullptr, "Slice input cannot be null");
StridedLayout lay(input->shape());
auto sliced_lay = lay.slice(slices_);
shape_ = sliced_lay.shape();
inputs_ = {std::move(input)};
}
// PadNode
PadNode::PadNode(std::shared_ptr<Node> operand, std::shared_ptr<Node> padding_value, std::vector<int64_t> edge_padding_low,
std::vector<int64_t> edge_padding_high, std::vector<int64_t> interior_padding)
: Node(NodeKind::kPad, operand->shape(), operand->Dtype(), operand->Sharding()),
low_(std::move(edge_padding_low)),
high_(std::move(edge_padding_high)),
interior_(std::move(interior_padding)) {
TENZR_CHECK(operand != nullptr && padding_value != nullptr, "Pad inputs cannot be null");
const size_t rank = operand->shape().Rank();
TENZR_CHECK(low_.size() == rank && high_.size() == rank, "Pad padding vectors must match operand rank");
interior_.resize(rank, 0);
std::vector<int64_t> out_dims(rank);
for (size_t i = 0; i < rank; ++i) {
const int64_t d = operand->shape().Dim(i);
const int64_t interior = i < interior_.size() ? interior_[i] : 0;
const int64_t space = d > 1 ? (d - 1) * interior : 0;
out_dims[i] = low_[i] + d + space + high_[i];
}
shape_ = Shape(out_dims);
inputs_ = {std::move(operand), std::move(padding_value)};
}
// GatherNode
GatherNode::GatherNode(std::shared_ptr<Node> operand, std::shared_ptr<Node> indices)
: Node(NodeKind::kGather, Shape({}), operand->Dtype(), operand->Sharding()) {
TENZR_CHECK(operand != nullptr && indices != nullptr, "Gather inputs cannot be null");
TENZR_CHECK(operand->shape().Rank() >= 1, "Gather operand must have rank >= 1");
const int64_t k = indices->shape().Dim(0);
std::vector<int64_t> out_dims = operand->shape().Dims();
if (!out_dims.empty()) {
out_dims[0] = k;
}
shape_ = Shape(out_dims);
inputs_ = {std::move(operand), std::move(indices)};
}
// ScatterNode
ScatterNode::ScatterNode(std::shared_ptr<Node> operand, std::shared_ptr<Node> indices, std::shared_ptr<Node> updates, std::string update_op)
: Node(NodeKind::kScatter, operand->shape(), operand->Dtype(), operand->Sharding()), update_op_(std::move(update_op)) {
TENZR_CHECK(operand != nullptr && indices != nullptr && updates != nullptr, "Scatter inputs cannot be null");
inputs_ = {std::move(operand), std::move(indices), std::move(updates)};
}
DynamicSliceNode::DynamicSliceNode(std::shared_ptr<Node> operand, std::vector<std::shared_ptr<Node>> start_indices, std::vector<int64_t> slice_sizes)
: Node(NodeKind::kDynamicSlice, Shape(slice_sizes), operand ? operand->Dtype() : DType::kFloat32,
operand ? operand->Sharding() : ShardingSpec::Replicated()),
slice_sizes_(std::move(slice_sizes)) {
TENZR_CHECK(operand != nullptr, "DynamicSliceNode operand cannot be null");
inputs_.push_back(std::move(operand));
for (auto& s : start_indices) {
inputs_.push_back(std::move(s));
}
}
DynamicUpdateSliceNode::DynamicUpdateSliceNode(std::shared_ptr<Node> operand, std::shared_ptr<Node> update, std::vector<std::shared_ptr<Node>> start_indices)
: Node(NodeKind::kDynamicUpdateSlice, operand ? operand->shape() : Shape{}, operand ? operand->Dtype() : DType::kFloat32,
operand ? operand->Sharding() : ShardingSpec::Replicated()) {
TENZR_CHECK(operand != nullptr && update != nullptr, "DynamicUpdateSliceNode inputs cannot be null");
inputs_.push_back(std::move(operand));
inputs_.push_back(std::move(update));
for (auto& s : start_indices) {
inputs_.push_back(std::move(s));
}
}
// ArgSortNode
ArgSortNode::ArgSortNode(std::shared_ptr<Node> values, std::shared_ptr<Node> indices, bool descending, int64_t axis, bool return_values)
: Node(NodeKind::kArgSort, values->shape(), return_values ? values->Dtype() : indices->Dtype(), values->Sharding()),
descending_(descending),
axis_(axis),
return_values_(return_values) {
TENZR_CHECK(values != nullptr && indices != nullptr, "ArgSort inputs cannot be null");
TENZR_CHECK(values->shape() == indices->shape(), "ArgSort values and indices must have the same shape");
inputs_ = {std::move(values), std::move(indices)};
}
// ConcatenateNode
ConcatenateNode::ConcatenateNode(std::vector<std::shared_ptr<Node>> inputs, int64_t axis)
: Node(NodeKind::kConcatenate, Shape({}), inputs.empty() ? DType::kFloat32 : inputs[0]->Dtype(),
inputs.empty() ? ShardingSpec::Replicated() : inputs[0]->Sharding()),
axis_(axis) {
TENZR_CHECK(!inputs.empty(), "Concatenate requires at least one input");
const size_t rank = inputs[0]->shape().Rank();
int64_t ax = axis < 0 ? axis + static_cast<int64_t>(rank) : axis;
TENZR_CHECK(ax >= 0 && ax < static_cast<int64_t>(rank), "Concatenate axis out of range");
std::vector<int64_t> out_dims = inputs[0]->shape().Dims();
int64_t sum_axis = 0;
for (const auto& in : inputs) {
TENZR_CHECK(in != nullptr, "Concatenate input cannot be null");
TENZR_CHECK(in->shape().Rank() == rank, "Concatenate inputs must all have the same rank");
sum_axis += in->shape().Dim(static_cast<size_t>(ax));
}
out_dims[static_cast<size_t>(ax)] = sum_axis;
shape_ = Shape(out_dims);
inputs_ = std::move(inputs);
}
// ReluNode
ReluNode::ReluNode(std::shared_ptr<Node> input) : Node(NodeKind::kRelu, input->shape(), input->Dtype(), input->Sharding()) {
TENZR_CHECK(input != nullptr, "Relu input cannot be null");
inputs_ = {std::move(input)};
}
// SigmoidNode
SigmoidNode::SigmoidNode(std::shared_ptr<Node> input) : Node(NodeKind::kSigmoid, input->shape(), input->Dtype(), input->Sharding()) {
TENZR_CHECK(input != nullptr, "Sigmoid input cannot be null");
inputs_ = {std::move(input)};
}
// CustomOpNode
CustomOpNode::CustomOpNode(std::string op_name, std::vector<std::shared_ptr<Node>> inputs, const OpDef* Definition)
: Node(NodeKind::kCustomOp, Shape{}, DType::kFloat32), op_name_str_(std::move(op_name)), op_def_(Definition) {
if (!op_def_) {
op_def_ = OpRegistry::Instance().Lookup(op_name_str_);
}
TENZR_CHECK(op_def_ != nullptr, "Unknown operator in CustomOpNode: " + op_name_str_);
inputs_ = std::move(inputs);
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_) {
TENZR_CHECK(in != nullptr, "CustomOpNode received null input node for op: " + op_name_str_);
in_shapes.push_back(in->shape());
in_dtypes.push_back(in->Dtype());
}
shape_ = op_def_->shape_fn ? op_def_->shape_fn(in_shapes) : (in_shapes.empty() ? Shape{} : in_shapes[0]);
dtype_ = op_def_->dtype_fn ? op_def_->dtype_fn(in_dtypes) : (in_dtypes.empty() ? DType::kFloat32 : in_dtypes[0]);
sharding_ = inputs_.empty() ? ShardingSpec::Replicated() : inputs_[0]->Sharding();
}
// ConvertNode
ConvertNode::ConvertNode(std::shared_ptr<Node> input, DType target_dtype) : Node(NodeKind::kConvert, input->shape(), target_dtype, input->Sharding()) {
TENZR_CHECK(input != nullptr, "ConvertNode input cannot be null");
inputs_ = {std::move(input)};
}
// ReduceNode
ReduceNode::ReduceNode(std::shared_ptr<Node> input, ReduceOp Reduction, std::vector<int64_t> dimensions, bool keepdims)
: Node(NodeKind::kReduce, Shape{}, input ? input->Dtype() : DType::kFloat32, input ? input->Sharding() : ShardingSpec::Replicated()),
reduce_op_(Reduction),
dimensions_(std::move(dimensions)),
keepdims_(keepdims) {
TENZR_CHECK(input != nullptr, "ReduceNode input cannot be null");
inputs_ = {std::move(input)};
const auto& in_shape = inputs_[0]->shape();
int64_t rank = static_cast<int64_t>(in_shape.Rank());
if (dimensions_.empty()) {
canonical_dimensions_.resize(static_cast<size_t>(rank));
std::iota(canonical_dimensions_.begin(), canonical_dimensions_.end(), 0);
} else {
std::vector<int64_t> canon;
canon.reserve(dimensions_.size());
for (int64_t d : dimensions_) {
int64_t c = (d < 0) ? (d + rank) : d;
TENZR_CHECK(c >= 0 && c < rank, "Dimension out of range: " + std::to_string(d) + " for rank " + std::to_string(rank));
canon.push_back(c);
}
std::sort(canon.begin(), canon.end());
canon.erase(std::unique(canon.begin(), canon.end()), canon.end());
canonical_dimensions_ = std::move(canon);
}
std::vector<int64_t> non_keep_dims;
std::vector<int64_t> keep_dims(static_cast<size_t>(rank));
for (int64_t i = 0; i < rank; ++i) {
bool is_reduced = std::find(canonical_dimensions_.begin(), canonical_dimensions_.end(), i) != canonical_dimensions_.end();
if (is_reduced) {
keep_dims[static_cast<size_t>(i)] = 1;
} else {
keep_dims[static_cast<size_t>(i)] = in_shape.Dim(static_cast<size_t>(i));
non_keep_dims.push_back(in_shape.Dim(static_cast<size_t>(i)));
}
}
non_keepdims_shape_ = Shape(non_keep_dims);
shape_ = keepdims_ ? Shape(keep_dims) : non_keepdims_shape_;
}
ReduceWindowNode::ReduceWindowNode(std::shared_ptr<Node> input, std::shared_ptr<Node> init_value, std::vector<int64_t> window_dimensions,
std::vector<int64_t> window_strides, std::vector<int64_t> padding_low, std::vector<int64_t> padding_high,
std::vector<int64_t> window_dilations, std::string Reduction)
: Node(NodeKind::kReduceWindow, Shape{}, input ? input->Dtype() : DType::kFloat32, input ? input->Sharding() : ShardingSpec::Replicated()),
window_dimensions_(std::move(window_dimensions)),
window_strides_(std::move(window_strides)),
padding_low_(std::move(padding_low)),
padding_high_(std::move(padding_high)),
window_dilations_(std::move(window_dilations)),
reduce_op_(std::move(Reduction)) {
TENZR_CHECK(input != nullptr && init_value != nullptr, "ReduceWindowNode inputs cannot be null");
inputs_ = {std::move(input), std::move(init_value)};
const auto& in_shape = inputs_[0]->shape();
const int64_t rank = static_cast<int64_t>(in_shape.Rank());
TENZR_CHECK(static_cast<int64_t>(window_dimensions_.size()) == rank, "ReduceWindow window_dimensions must match the operand rank");
if (window_strides_.empty()) {
window_strides_.assign(static_cast<size_t>(rank), 1);
}
if (padding_low_.empty()) {
padding_low_.assign(static_cast<size_t>(rank), 0);
}
if (padding_high_.empty()) {
padding_high_.assign(static_cast<size_t>(rank), 0);
}
if (window_dilations_.empty()) {
window_dilations_.assign(static_cast<size_t>(rank), 1);
}
std::vector<int64_t> out_dims(static_cast<size_t>(rank));
for (int64_t d = 0; d < rank; ++d) {
const int64_t in_d = in_shape.Dim(static_cast<size_t>(d));
const int64_t effective = (window_dimensions_[static_cast<size_t>(d)] - 1) * window_dilations_[static_cast<size_t>(d)] + 1;
const int64_t windowed = in_d + padding_low_[static_cast<size_t>(d)] + padding_high_[static_cast<size_t>(d)] - effective;
TENZR_CHECK(windowed >= 0, "ReduceWindow window is larger than the padded operand");
out_dims[static_cast<size_t>(d)] = windowed / window_strides_[static_cast<size_t>(d)] + 1;
}
shape_ = Shape(out_dims);
}
std::string ReduceNode::OpName() const {
switch (reduce_op_) {
case ReduceOp::kSum:
return "ReduceSum";
case ReduceOp::kProd:
return "ReduceProd";
case ReduceOp::kMax:
return "ReduceMax";
case ReduceOp::kMin:
return "ReduceMin";
}
return "Reduce";
}
Shape ReduceNode::ExpandedShape() const {
const auto& in_shape = inputs_[0]->shape();
int64_t rank = static_cast<int64_t>(in_shape.Rank());
std::vector<int64_t> exp_dims(static_cast<size_t>(rank));
for (int64_t i = 0; i < rank; ++i) {
bool is_reduced = std::find(canonical_dimensions_.begin(), canonical_dimensions_.end(), i) != canonical_dimensions_.end();
exp_dims[static_cast<size_t>(i)] = is_reduced ? 1 : in_shape.Dim(static_cast<size_t>(i));
}
return Shape(exp_dims);
}
} // namespace tenzr