// 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/plax/plax.h"
#include <algorithm>
#include <cstring>
#include <stdexcept>
#include <string>
#include <vector>
#include "tenzr/common/error.h"
#include "tenzr/common/logging.h"
#include "tenzr/frontend/dtype.h"
#include "tenzr/frontend/memory_pool.h"
#include "tenzr/sharding/named_sharding.h"
#include "tenzr/txpr/custom_vjp.h"
#include "tenzr/txpr/tracer.h"
#include "tenzr/txpr/txpr.h"
namespace tenzr::plax {
namespace {
inline void CheckStrictBinary(const Array& lhs, const Array& rhs, const char* op_name) {
if (lhs.Dtype() != rhs.Dtype()) {
TENZR_THROW(std::string("plax::") + op_name + " requires strictly identical dtypes, but got " + std::string(tenzr::ToString(lhs.Dtype())) + " vs " +
std::string(tenzr::ToString(rhs.Dtype())) + ". Use plax::convert_element_type to explicitly cast.");
}
if (lhs.shape() != rhs.shape()) {
TENZR_THROW(std::string("plax::") + op_name + " requires strictly identical shapes, but got " + lhs.shape().ToString() + " vs " +
rhs.shape().ToString() + ". Use plax::broadcast_in_dim for explicit broadcasting.");
}
}
} // namespace
// Strict Elementwise Binary Primitives
Array Add(const Array& lhs, const Array& rhs) {
CheckStrictBinary(lhs, rhs, "add");
return lhs + rhs;
}
Array Sub(const Array& lhs, const Array& rhs) {
CheckStrictBinary(lhs, rhs, "sub");
return lhs - rhs;
}
Array Mul(const Array& lhs, const Array& rhs) {
CheckStrictBinary(lhs, rhs, "mul");
return lhs * rhs;
}
Array Div(const Array& lhs, const Array& rhs) {
CheckStrictBinary(lhs, rhs, "div");
return lhs / rhs;
}
Array Pow(const Array& lhs, const Array& rhs) {
CheckStrictBinary(lhs, rhs, "pow");
return lhs.Pow(rhs);
}
Array Maximum(const Array& lhs, const Array& rhs) {
CheckStrictBinary(lhs, rhs, "maximum");
return lhs.Maximum(rhs);
}
Array Minimum(const Array& lhs, const Array& rhs) {
CheckStrictBinary(lhs, rhs, "minimum");
return lhs.Minimum(rhs);
}
Array Remainder(const Array& lhs, const Array& rhs) {
CheckStrictBinary(lhs, rhs, "remainder");
return lhs.Remainder(rhs);
}
Array Atan2(const Array& lhs, const Array& rhs) {
CheckStrictBinary(lhs, rhs, "atan2");
return lhs.Atan2(rhs);
}
// Strict Comparison Primitives (returns Bool array)
Array Eq(const Array& lhs, const Array& rhs) {
CheckStrictBinary(lhs, rhs, "eq");
return lhs.Equal(rhs);
}
Array Ne(const Array& lhs, const Array& rhs) {
CheckStrictBinary(lhs, rhs, "ne");
return lhs.NotEqual(rhs);
}
Array Lt(const Array& lhs, const Array& rhs) {
CheckStrictBinary(lhs, rhs, "lt");
return lhs.LessThan(rhs);
}
Array Le(const Array& lhs, const Array& rhs) {
CheckStrictBinary(lhs, rhs, "le");
return lhs.LessEqual(rhs);
}
Array Gt(const Array& lhs, const Array& rhs) {
CheckStrictBinary(lhs, rhs, "gt");
return lhs.GreaterThan(rhs);
}
Array Ge(const Array& lhs, const Array& rhs) {
CheckStrictBinary(lhs, rhs, "ge");
return lhs.GreaterEqual(rhs);
}
// Strict Elementwise Unary Primitives
Array Neg(const Array& x) { return x.Neg(); }
Array Exp(const Array& x) { return x.Exp(); }
Array Log(const Array& x) { return x.Log(); }
Array tanh(const Array& x) { return x.tanh(); }
Array Abs(const Array& x) { return x.Abs(); }
Array Sqrt(const Array& x) { return x.Sqrt(); }
Array Rsqrt(const Array& x) { return x.Rsqrt(); }
Array Sin(const Array& x) { return x.Sin(); }
Array Cos(const Array& x) { return x.Cos(); }
Array Tan(const Array& x) { return x.Tan(); }
Array Ceil(const Array& x) { return x.Ceil(); }
Array Floor(const Array& x) { return x.Floor(); }
Array Round(const Array& x) { return x.Round(); }
Array Sign(const Array& x) { return x.Sign(); }
// Reductions
Array ReduceSum(const Array& operand, const std::vector<int64_t>& axes, bool keepdims) {
if (axes.empty()) {
return operand.Sum(std::nullopt, keepdims);
}
return operand.Sum(axes, keepdims);
}
Array ReduceMax(const Array& operand, const std::vector<int64_t>& axes, bool keepdims) {
if (axes.empty()) {
return operand.Max(std::nullopt, keepdims);
}
return operand.Max(axes, keepdims);
}
Array ReduceMin(const Array& operand, const std::vector<int64_t>& axes, bool keepdims) {
if (axes.empty()) {
return operand.Min(std::nullopt, keepdims);
}
return operand.Min(axes, keepdims);
}
Array ReduceProd(const Array& operand, const std::vector<int64_t>& axes, bool keepdims) {
if (axes.empty()) {
return operand.Prod(std::nullopt, keepdims);
}
return operand.Prod(axes, keepdims);
}
Array ReduceWindow(const Array& operand, const Array& init_value, const std::vector<int64_t>& window_dimensions, const std::vector<int64_t>& window_strides,
const std::vector<int64_t>& padding_low, const std::vector<int64_t>& padding_high, std::string_view reducer) {
const int64_t Rank = static_cast<int64_t>(operand.shape().Rank());
TENZR_CHECK(static_cast<int64_t>(window_dimensions.size()) == Rank, "plax::reduce_window window_dimensions must match the operand rank");
TENZR_CHECK(reducer == "add" || reducer == "max" || reducer == "min" || reducer == "multiply", "plax::reduce_window unknown reducer");
const std::vector<int64_t> strides = window_strides.empty() ? std::vector<int64_t>(static_cast<size_t>(Rank), 1) : window_strides;
const std::vector<int64_t> low = padding_low.empty() ? std::vector<int64_t>(static_cast<size_t>(Rank), 0) : padding_low;
const std::vector<int64_t> high = padding_high.empty() ? std::vector<int64_t>(static_cast<size_t>(Rank), 0) : padding_high;
const std::vector<int64_t> dilations(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 = operand.shape().Dim(static_cast<size_t>(d));
const int64_t effective = (window_dimensions[static_cast<size_t>(d)] - 1) * dilations[static_cast<size_t>(d)] + 1;
const int64_t windowed = in_d + low[static_cast<size_t>(d)] + high[static_cast<size_t>(d)] - effective;
TENZR_CHECK(windowed >= 0, "plax::reduce_window window is larger than the padded operand");
out_dims[static_cast<size_t>(d)] = windowed / strides[static_cast<size_t>(d)] + 1;
}
Shape out_shape(out_dims);
if (txpr::TracerContext::IsActive() || operand.IsSymbolic() || init_value.IsSymbolic()) {
auto* tape = txpr::TracerContext::CurrentTxpr();
TENZR_CHECK(tape != nullptr, "Active tape required for symbolic reduce_window");
txpr::VarId in_id = operand.ObtainVarId(tape);
txpr::VarId init_id = init_value.ObtainVarId(tape);
txpr::AttributePayload attrs;
attrs.window_dimensions = window_dimensions;
attrs.window_strides = strides;
attrs.base_dilations.assign(static_cast<size_t>(Rank), 1);
attrs.window_dilations = dilations;
attrs.edge_padding_low = low;
attrs.edge_padding_high = high;
attrs.reduce_window_op = std::string(reducer);
txpr::VarId out_id = tape->AppendInstruction(txpr::OpCode::kReduceWindow, {in_id, init_id}, out_shape, operand.Dtype(), std::move(attrs));
return Array::Symbolic(out_id, out_shape, operand.Dtype());
}
return Array(std::make_shared<ReduceWindowNode>(operand.node(), init_value.node(), window_dimensions, strides, low, high, dilations, std::string(reducer)));
}
// Windowing, Slicing & Padding
Array slice(const Array& operand, const std::vector<int64_t>& start_indices, const std::vector<int64_t>& limit_indices, const std::vector<int64_t>& strides) {
size_t Rank = operand.shape().Rank();
if (start_indices.size() != Rank || limit_indices.size() != Rank) {
TENZR_THROW("plax::slice indices length must match operand rank");
}
std::vector<Slice> slices;
slices.reserve(Rank);
for (size_t i = 0; i < Rank; ++i) {
int64_t st = start_indices[i];
int64_t sp = limit_indices[i];
int64_t step = (i < strides.size()) ? strides[i] : 1;
slices.push_back(Slice(st, sp, step));
}
return operand.slice(slices);
}
Array Pad(const Array& operand, const Array& padding_value, const std::vector<int64_t>& edge_padding_low, const std::vector<int64_t>& edge_padding_high,
const std::vector<int64_t>& interior_padding) {
if (operand.Dtype() != padding_value.Dtype()) {
TENZR_THROW("plax::pad operand and padding_value dtypes must match");
}
size_t rank = operand.shape().Rank();
if (edge_padding_low.size() != rank || edge_padding_high.size() != rank) {
TENZR_THROW("plax::pad edge_padding vector sizes must match operand rank");
}
std::vector<int64_t> out_dims(rank);
for (size_t i = 0; i < rank; ++i) {
int64_t d = operand.shape().Dim(i);
int64_t int_pad = (i < interior_padding.size()) ? interior_padding[i] : 0;
int64_t interior_space = (d > 1) ? ((d - 1) * int_pad) : 0;
out_dims[i] = edge_padding_low[i] + d + interior_space + edge_padding_high[i];
}
Shape out_shape(std::move(out_dims));
if (txpr::TracerContext::IsActive() || operand.IsSymbolic() || padding_value.IsSymbolic()) {
auto* tape = txpr::TracerContext::CurrentTxpr();
TENZR_CHECK(tape != nullptr, "Active tape required for symbolic pad");
TENZR_CHECK(padding_value.shape().NumElements() == 1, "plax::pad padding_value must be a scalar");
txpr::AttributePayload attrs;
attrs.edge_padding_low = edge_padding_low;
attrs.edge_padding_high = edge_padding_high;
attrs.interior_padding.resize(rank, 0);
for (size_t i = 0; i < rank && i < interior_padding.size(); ++i) {
attrs.interior_padding[i] = interior_padding[i];
}
txpr::VarId in_id = operand.ObtainVarId(tape);
txpr::VarId pv_id = padding_value.ObtainVarId(tape);
txpr::VarId out_id = tape->AppendInstruction(txpr::OpCode::kPad, {in_id, pv_id}, out_shape, operand.Dtype(), std::move(attrs));
return Array::Symbolic(out_id, out_shape, operand.Dtype());
}
return Array(std::make_shared<PadNode>(operand.node(), padding_value.node(), edge_padding_low, edge_padding_high, interior_padding));
}
// Indexing
Array Gather(const Array& operand, const Array& start_indices, int64_t axis) {
int64_t Rank = static_cast<int64_t>(operand.shape().Rank());
if (axis < 0) {
axis += Rank;
}
if (axis < 0 || axis >= Rank) {
TENZR_THROW("plax::gather axis out of range");
}
if (start_indices.shape().Rank() != 1) {
TENZR_THROW("plax::gather start_indices must be 1-D");
}
int64_t k = start_indices.shape().Dim(0);
if (txpr::TracerContext::IsActive() || operand.IsSymbolic() || start_indices.IsSymbolic()) {
auto* tape = txpr::TracerContext::CurrentTxpr();
TENZR_CHECK(tape != nullptr, "Active tape required for symbolic gather");
txpr::VarId op_id = operand.ObtainVarId(tape);
txpr::VarId idx_id = start_indices.ObtainVarId(tape);
txpr::VarId idx2 = tape->AppendInstruction(txpr::OpCode::kReshape, {idx_id}, Shape({k, 1}), start_indices.Dtype());
// stablehlo.gather produces [k] + (operand dims except axis); reorder for axis != 0.
std::vector<int64_t> nat_dims;
nat_dims.reserve(static_cast<size_t>(Rank));
nat_dims.push_back(k);
for (int64_t d = 0; d < Rank; ++d) {
if (d != axis) {
nat_dims.push_back(operand.shape().Dim(static_cast<size_t>(d)));
}
}
Shape nat_shape(std::move(nat_dims));
txpr::AttributePayload attrs;
attrs.gather_axis = axis;
txpr::VarId out_id = tape->AppendInstruction(txpr::OpCode::kGather, {op_id, idx2}, nat_shape, operand.Dtype(), std::move(attrs));
Array nat = Array::Symbolic(out_id, nat_shape, operand.Dtype());
if (axis == 0) {
return nat;
}
std::vector<int64_t> perm;
perm.reserve(static_cast<size_t>(Rank));
for (int64_t d = 1; d <= axis; ++d) {
perm.push_back(d);
}
perm.push_back(0);
for (int64_t d = axis + 1; d < Rank; ++d) {
perm.push_back(d);
}
return nat.Permute(perm);
}
if (axis == 0) {
Array idx2 = start_indices.Reshape(Shape({k, 1}));
return Array(std::make_shared<GatherNode>(operand.node(), idx2.node()));
}
// axis != 0: move the gathered axis to the front, gather on axis 0 (the only
// native eager gather), then restore the natural dimension order. This keeps
// the operation traceable/differentiable/dtype-generic instead of a host loop.
std::vector<int64_t> to_front;
to_front.reserve(static_cast<size_t>(Rank));
to_front.push_back(axis);
for (int64_t d = 0; d < Rank; ++d) {
if (d != axis) {
to_front.push_back(d);
}
}
std::vector<int64_t> restore;
restore.reserve(static_cast<size_t>(Rank));
for (int64_t d = 1; d <= axis; ++d) {
restore.push_back(d);
}
restore.push_back(0);
for (int64_t d = axis + 1; d < Rank; ++d) {
restore.push_back(d);
}
Array front = operand.Permute(to_front);
Array idx2 = start_indices.Reshape(Shape({k, 1}));
Array gathered = Array(std::make_shared<GatherNode>(front.node(), idx2.node()));
return gathered.Permute(restore);
}
namespace {
std::vector<Array> SplitStartIndices(const Array& start_indices, size_t rank) {
TENZR_CHECK(start_indices.shape().NumElements() == static_cast<int64_t>(rank), "plax dynamic slice start_indices must have one entry per dimension");
std::vector<Array> starts;
starts.reserve(rank);
for (size_t d = 0; d < rank; ++d) {
starts.push_back(start_indices.slice({Slice(static_cast<int64_t>(d), static_cast<int64_t>(d) + 1, 1)}).Reshape(Shape({})));
}
return starts;
}
} // namespace
Array DynamicSlice(const Array& operand, const std::vector<Array>& start_indices, const Shape& slice_sizes) {
TENZR_CHECK(start_indices.size() == operand.shape().Rank(), "plax::dynamic_slice requires one start index per operand dimension");
if (txpr::TracerContext::IsActive() || operand.IsSymbolic()) {
auto* tape = txpr::TracerContext::CurrentTxpr();
TENZR_CHECK(tape != nullptr, "Active tape required for symbolic dynamic_slice");
std::vector<txpr::VarId> ins;
ins.reserve(start_indices.size() + 1);
ins.push_back(operand.ObtainVarId(tape));
for (const auto& s : start_indices) {
ins.push_back(s.ObtainVarId(tape));
}
txpr::AttributePayload attrs;
attrs.slice_sizes = slice_sizes.Dims();
txpr::VarId out = tape->AppendInstruction(txpr::OpCode::kDynamicSlice, std::move(ins), slice_sizes, operand.Dtype(), std::move(attrs));
return Array::Symbolic(out, slice_sizes, operand.Dtype());
}
std::vector<std::shared_ptr<Node>> starts;
starts.reserve(start_indices.size());
for (const auto& s : start_indices) {
starts.push_back(s.node());
}
return Array(std::make_shared<DynamicSliceNode>(operand.node(), std::move(starts), slice_sizes.Dims()));
}
Array DynamicSlice(const Array& operand, const Array& start_indices, const std::vector<int64_t>& slice_sizes) {
return DynamicSlice(operand, SplitStartIndices(start_indices, operand.shape().Rank()), Shape(slice_sizes));
}
Array DynamicUpdateSlice(const Array& operand, const Array& update, const std::vector<Array>& start_indices) {
TENZR_CHECK(start_indices.size() == operand.shape().Rank(), "plax::dynamic_update_slice requires one start index per operand dimension");
if (txpr::TracerContext::IsActive() || operand.IsSymbolic() || update.IsSymbolic()) {
auto* tape = txpr::TracerContext::CurrentTxpr();
TENZR_CHECK(tape != nullptr, "Active tape required for symbolic dynamic_update_slice");
std::vector<txpr::VarId> ins;
ins.reserve(start_indices.size() + 2);
ins.push_back(operand.ObtainVarId(tape));
ins.push_back(update.ObtainVarId(tape));
for (const auto& s : start_indices) {
ins.push_back(s.ObtainVarId(tape));
}
txpr::VarId out = tape->AppendInstruction(txpr::OpCode::kDynamicUpdateSlice, std::move(ins), operand.shape(), operand.Dtype());
return Array::Symbolic(out, operand.shape(), operand.Dtype());
}
std::vector<std::shared_ptr<Node>> starts;
starts.reserve(start_indices.size());
for (const auto& s : start_indices) {
starts.push_back(s.node());
}
return Array(std::make_shared<DynamicUpdateSliceNode>(operand.node(), update.node(), std::move(starts)));
}
Array DynamicUpdateSlice(const Array& operand, const Array& update, const Array& start_indices) {
return DynamicUpdateSlice(operand, update, SplitStartIndices(start_indices, operand.shape().Rank()));
}
namespace {
std::vector<int32_t> MakeIotaData(const Shape& shape, int64_t axis) {
const int64_t Rank = static_cast<int64_t>(shape.Rank());
int64_t inner = 1;
for (int64_t d = axis + 1; d < Rank; ++d) {
inner *= shape.Dim(static_cast<size_t>(d));
}
const int64_t dim = shape.Dim(static_cast<size_t>(axis));
const int64_t total = shape.NumElements();
const int64_t outer = total / (dim * inner);
std::vector<int32_t> v(static_cast<size_t>(total));
for (int64_t o = 0; o < outer; ++o) {
for (int64_t i = 0; i < dim; ++i) {
for (int64_t in = 0; in < inner; ++in) {
v[static_cast<size_t>((o * dim + i) * inner + in)] = static_cast<int32_t>(i);
}
}
}
return v;
}
Array MakeInt32Constant(const std::vector<int32_t>& data, const Shape& shape) {
if (txpr::TracerContext::IsActive()) {
auto* tape = txpr::TracerContext::CurrentTxpr();
if (tape != nullptr) {
txpr::AttributePayload attrs;
attrs.int_tensor_data.assign(data.begin(), data.end());
txpr::VarId v = tape->AppendInstruction(txpr::OpCode::kConstant, {}, shape, DType::kInt32, std::move(attrs));
return Array::Symbolic(v, shape, DType::kInt32);
}
}
Array a = Array::Zeros(shape, DType::kInt32);
a.SetData<int32_t>(data);
return a;
}
// Native index vector: stablehlo.iota along axis when tracing, else a host
// iota. Avoids serialising a full-shape dense constant into the emitted MLIR.
Array MakeIotaArray(const Shape& shape, int64_t axis) {
if (txpr::TracerContext::IsActive()) {
auto* tape = txpr::TracerContext::CurrentTxpr();
if (tape != nullptr) {
txpr::AttributePayload attrs;
attrs.iota_dimension = axis;
txpr::VarId v = tape->AppendInstruction(txpr::OpCode::kConstant, {}, shape, DType::kInt32, std::move(attrs));
return Array::Symbolic(v, shape, DType::kInt32);
}
}
Array a = Array::Zeros(shape, DType::kInt32);
a.SetData<int32_t>(MakeIotaData(shape, axis));
return a;
}
Array SortImpl(const Array& operand, int64_t axis, bool descending, bool return_values) {
int64_t Rank = static_cast<int64_t>(operand.shape().Rank());
TENZR_CHECK(Rank >= 1, "plax sort requires rank >= 1");
if (axis < 0) {
axis += Rank;
}
TENZR_CHECK(axis >= 0 && axis < Rank, "plax sort axis out of range");
const Shape& shape = operand.shape();
const DType out_dtype = return_values ? operand.Dtype() : DType::kInt32;
if (txpr::TracerContext::IsActive() || operand.IsSymbolic()) {
auto* tape = txpr::TracerContext::CurrentTxpr();
TENZR_CHECK(tape != nullptr, "Active tape required for symbolic sort");
txpr::VarId v_id = operand.ObtainVarId(tape);
txpr::AttributePayload attrs;
attrs.sort_dimension = axis;
attrs.sort_descending = descending;
attrs.sort_return_values = return_values;
if (return_values) {
// Values-only sort: single-operand stablehlo.sort, no index payload.
txpr::VarId out_id = tape->AppendInstruction(txpr::OpCode::kArgSort, {v_id}, shape, out_dtype, std::move(attrs));
return Array::Symbolic(out_id, shape, out_dtype);
}
Array iota_arr = MakeIotaArray(shape, axis);
txpr::VarId iota_id = iota_arr.ObtainVarId(tape);
txpr::VarId out_id = tape->AppendInstruction(txpr::OpCode::kArgSort, {v_id, iota_id}, shape, out_dtype, std::move(attrs));
return Array::Symbolic(out_id, shape, out_dtype);
}
if (return_values) {
// The index payload is unused for a values-only sort; reuse the operand
// as the placeholder input to avoid materialising a host iota.
return Array(std::make_shared<ArgSortNode>(operand.node(), operand.node(), descending, axis, true));
}
Array iota_arr = Array::Zeros(shape, DType::kInt32);
iota_arr.SetData<int32_t>(MakeIotaData(shape, axis));
return Array(std::make_shared<ArgSortNode>(operand.node(), iota_arr.node(), descending, axis, false));
}
} // namespace
Array Scatter(const Array& operand, const Array& indices, const Array& updates, int64_t axis, std::string_view update) {
int64_t Rank = static_cast<int64_t>(operand.shape().Rank());
if (axis < 0) {
axis += Rank;
}
if (axis < 0 || axis >= Rank) {
TENZR_THROW("plax::scatter axis out of range");
}
if (operand.Dtype() != updates.Dtype()) {
TENZR_THROW("plax::scatter operand and updates dtypes must match");
}
if (update != "add" && update != "set" && update != "multiply" && update != "min" && update != "max") {
TENZR_THROW("plax::scatter unknown update mode: " + std::string(update));
}
int64_t k = indices.shape().Dim(0);
if (txpr::TracerContext::IsActive() || operand.IsSymbolic() || indices.IsSymbolic() || updates.IsSymbolic()) {
auto* tape = txpr::TracerContext::CurrentTxpr();
TENZR_CHECK(tape != nullptr, "Active tape required for symbolic scatter");
txpr::VarId op_id = operand.ObtainVarId(tape);
txpr::VarId idx_id = indices.ObtainVarId(tape);
txpr::VarId upd_id = updates.ObtainVarId(tape);
txpr::VarId idx2 = tape->AppendInstruction(txpr::OpCode::kReshape, {idx_id}, Shape({k, 1}), indices.Dtype());
txpr::AttributePayload attrs;
attrs.scatter_axis = axis;
attrs.scatter_update = std::string(update);
txpr::VarId out_id = tape->AppendInstruction(txpr::OpCode::kScatter, {op_id, idx2, upd_id}, operand.shape(), operand.Dtype(), std::move(attrs));
return Array::Symbolic(out_id, operand.shape(), operand.Dtype());
}
if (axis == 0) {
Array idx2 = indices.Reshape(Shape({k, 1}));
return Array(std::make_shared<ScatterNode>(operand.node(), idx2.node(), updates.node(), std::string(update)));
}
// axis != 0: move the scatter axis to the front, scatter on axis 0 (the only
// native eager scatter), then restore the natural dimension order. Traceable,
// differentiable and dtype-generic instead of a float32-only host loop.
TENZR_CHECK(static_cast<int64_t>(updates.shape().Rank()) == Rank, "plax::scatter updates must have the same rank as the operand");
std::vector<int64_t> to_front;
to_front.reserve(static_cast<size_t>(Rank));
to_front.push_back(axis);
for (int64_t d = 0; d < Rank; ++d) {
if (d != axis) {
to_front.push_back(d);
}
}
std::vector<int64_t> restore;
restore.reserve(static_cast<size_t>(Rank));
for (int64_t d = 1; d <= axis; ++d) {
restore.push_back(d);
}
restore.push_back(0);
for (int64_t d = axis + 1; d < Rank; ++d) {
restore.push_back(d);
}
Array op_front = operand.Permute(to_front);
Array up_front = updates.Permute(to_front);
Array idx2 = indices.Reshape(Shape({k, 1}));
Array scattered = Array(std::make_shared<ScatterNode>(op_front.node(), idx2.node(), up_front.node(), std::string(update)));
return scattered.Permute(restore);
}
Array ScatterAdd(const Array& operand, const Array& indices, const Array& updates, int64_t axis) { return Scatter(operand, indices, updates, axis, "add"); }
Array Argsort(const Array& operand, int64_t axis, bool descending) { return SortImpl(operand, axis, descending, false); }
Array Sort(const Array& operand, int64_t axis, bool descending) { return SortImpl(operand, axis, descending, true); }
namespace {
// argmax/argmin as an O(n) masked reduction: the first position attaining
// the extreme is min{ iota | operand == extreme } (ties -> lowest index,
// matching NumPy). Composed from primitives so it stays traceable/vmap-able.
Array ArgExtreme(const Array& operand, int64_t axis, bool find_max) {
const int64_t rank = static_cast<int64_t>(operand.shape().Rank());
TENZR_CHECK(rank >= 1, "plax argmax/argmin requires rank >= 1");
if (axis < 0) {
axis += rank;
}
TENZR_CHECK(axis >= 0 && axis < rank, "plax argmax/argmin axis out of range");
const Shape& shape = operand.shape();
const std::vector<int64_t> axes{axis};
Array extreme = find_max ? ReduceMax(operand, axes, true) : ReduceMin(operand, axes, true);
Array is_extreme = operand.Equal(extreme);
Array positions = MakeIotaArray(shape, axis);
Array sentinel = MakeInt32Constant({static_cast<int32_t>(shape.Dim(static_cast<size_t>(axis)))}, Shape({}));
Array candidates = Array::CustomOp("select", {is_extreme, positions, sentinel});
return ReduceMin(candidates, axes, false);
}
} // namespace
Array Argmax(const Array& operand, int64_t axis) { return ArgExtreme(operand, axis, true); }
Array Argmin(const Array& operand, int64_t axis) { return ArgExtreme(operand, axis, false); }
Array Concatenate(const std::vector<Array>& arrays, int64_t axis) {
TENZR_CHECK(!arrays.empty(), "plax::concatenate requires at least one array");
if (arrays.size() == 1) {
return arrays[0];
}
const size_t Rank = arrays[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), "plax::concatenate axis out of range");
if (txpr::TracerContext::IsActive() || arrays[0].IsSymbolic()) {
auto* tape = txpr::TracerContext::CurrentTxpr();
TENZR_CHECK(tape != nullptr, "Active tape required for symbolic concatenate");
std::vector<txpr::VarId> ids;
ids.reserve(arrays.size());
std::vector<int64_t> out_dims = arrays[0].shape().Dims();
int64_t sum_axis = 0;
for (const auto& a : arrays) {
TENZR_CHECK(a.Dtype() == arrays[0].Dtype(), "plax::concatenate requires identical dtypes");
TENZR_CHECK(a.shape().Rank() == Rank, "plax::concatenate requires equal ranks");
ids.push_back(a.ObtainVarId(tape));
sum_axis += a.shape().Dim(static_cast<size_t>(ax));
}
out_dims[static_cast<size_t>(ax)] = sum_axis;
Shape out_shape(out_dims);
txpr::AttributePayload attrs;
attrs.concat_axis = ax;
txpr::VarId out = tape->AppendInstruction(txpr::OpCode::kConcatenate, std::move(ids), out_shape, arrays[0].Dtype(), std::move(attrs));
return Array::Symbolic(out, out_shape, arrays[0].Dtype());
}
std::vector<std::shared_ptr<Node>> nodes;
nodes.reserve(arrays.size());
for (const auto& a : arrays) {
TENZR_CHECK(a.Dtype() == arrays[0].Dtype(), "plax::concatenate requires identical dtypes");
nodes.push_back(a.node());
}
return Array(std::make_shared<ConcatenateNode>(std::move(nodes), ax));
}
Array Stack(const std::vector<Array>& arrays, int64_t axis) {
std::vector<Array> unsq;
unsq.reserve(arrays.size());
for (const auto& a : arrays) {
unsq.push_back(a.Unsqueeze(axis));
}
return Concatenate(unsq, axis);
}
Array TakeAlongAxis(const Array& operand, const Array& indices, int64_t axis) {
const int64_t Rank = static_cast<int64_t>(operand.shape().Rank());
TENZR_CHECK(Rank >= 1, "plax::take_along_axis requires rank >= 1");
TENZR_CHECK(static_cast<int64_t>(indices.shape().Rank()) == Rank, "plax::take_along_axis indices must have the same rank as the operand");
if (axis < 0) {
axis += Rank;
}
TENZR_CHECK(axis >= 0 && axis < Rank, "plax::take_along_axis axis out of range");
const Shape& shape = operand.shape();
// NumPy semantics: the result shape is the broadcast of indices and operand,
// with the axis dimension taken from indices (so k < operand_axis_dim is
// allowed). out[u] = operand[u with axis := idx[u]].
std::vector<int64_t> out_dims(static_cast<size_t>(Rank));
for (int64_t d = 0; d < Rank; ++d) {
const int64_t od = shape.Dim(static_cast<size_t>(d));
const int64_t id = indices.shape().Dim(static_cast<size_t>(d));
if (d == axis) {
out_dims[static_cast<size_t>(d)] = id;
} else {
TENZR_CHECK(od == id || od == 1 || id == 1, "plax::take_along_axis operand and indices are not broadcast-compatible");
out_dims[static_cast<size_t>(d)] = std::max(od, id);
}
}
Shape out_shape(out_dims);
const int64_t total_operand = shape.NumElements();
const int64_t total_out = out_shape.NumElements();
std::vector<int64_t> operand_stride(static_cast<size_t>(Rank), 1);
std::vector<int64_t> out_stride(static_cast<size_t>(Rank), 1);
for (int64_t d = Rank - 2; d >= 0; --d) {
operand_stride[static_cast<size_t>(d)] = operand_stride[static_cast<size_t>(d) + 1] * shape.Dim(static_cast<size_t>(d) + 1);
out_stride[static_cast<size_t>(d)] = out_stride[static_cast<size_t>(d) + 1] * out_shape.Dim(static_cast<size_t>(d) + 1);
}
const int64_t axis_stride = operand_stride[static_cast<size_t>(axis)];
// The flat operand index is base[u] + idx[u] * axis_stride, where base[u] is the
// operand flat index with the axis coordinate zeroed. Lowering to a single 1-D
// stablehlo.gather keeps this O(numel) instead of an O(numel * axis_dim) one-hot.
std::vector<int32_t> base(static_cast<size_t>(total_out));
for (int64_t f = 0; f < total_out; ++f) {
int64_t flat = 0;
for (int64_t d = 0; d < Rank; ++d) {
if (d == axis) {
continue;
}
int64_t coord = (f / out_stride[static_cast<size_t>(d)]) % out_shape.Dim(static_cast<size_t>(d));
if (shape.Dim(static_cast<size_t>(d)) == 1) {
coord = 0;
}
flat += coord * operand_stride[static_cast<size_t>(d)];
}
base[static_cast<size_t>(f)] = static_cast<int32_t>(flat);
}
Array idx = indices.Astype(DType::kInt32);
if (idx.shape() != out_shape) {
idx = idx.BroadcastTo(out_shape);
}
Array base_arr = MakeInt32Constant(base, out_shape);
Array stride_arr = MakeInt32Constant({static_cast<int32_t>(axis_stride)}, Shape({}));
Array linear = idx * stride_arr + base_arr;
Array flat_operand = operand.Reshape(Shape({total_operand}));
Array flat_linear = linear.Reshape(Shape({total_out}));
return Gather(flat_operand, flat_linear, 0).Reshape(out_shape);
}
Array PutAlongAxis(const Array& operand, const Array& indices, const Array& values, int64_t axis, std::string_view update) {
const int64_t Rank = static_cast<int64_t>(operand.shape().Rank());
TENZR_CHECK(Rank >= 1, "plax::put_along_axis requires rank >= 1");
TENZR_CHECK(static_cast<int64_t>(indices.shape().Rank()) == Rank, "plax::put_along_axis indices must have the same rank as the operand");
TENZR_CHECK(static_cast<int64_t>(values.shape().Rank()) == Rank, "plax::put_along_axis values must have the same rank as the operand");
TENZR_CHECK(operand.Dtype() == values.Dtype(), "plax::put_along_axis operand and values dtypes must match");
if (axis < 0) {
axis += Rank;
}
TENZR_CHECK(axis >= 0 && axis < Rank, "plax::put_along_axis axis out of range");
const Shape& shape = operand.shape();
// indices/values define the update shape (equal to the operand shape except
// along axis, where it is the number of updates). NumPy broadcasts them.
Shape update_shape = Shape::Broadcast(indices.shape(), values.shape());
TENZR_CHECK(static_cast<int64_t>(update_shape.Rank()) == Rank, "plax::put_along_axis indices/values must have the same rank as the operand");
Array idx = indices.Astype(DType::kInt32);
if (idx.shape() != update_shape) {
idx = idx.BroadcastTo(update_shape);
}
Array vals = values;
if (vals.shape() != update_shape) {
vals = vals.BroadcastTo(update_shape);
}
// Operand and update row-major strides; the destination flat index of an update
// at (outer, j, inner) is base + idx[j] * operand_axis_stride.
std::vector<int64_t> operand_stride(static_cast<size_t>(Rank), 1);
std::vector<int64_t> update_stride(static_cast<size_t>(Rank), 1);
for (int64_t d = Rank - 2; d >= 0; --d) {
operand_stride[static_cast<size_t>(d)] = operand_stride[static_cast<size_t>(d) + 1] * shape.Dim(static_cast<size_t>(d) + 1);
update_stride[static_cast<size_t>(d)] = update_stride[static_cast<size_t>(d) + 1] * update_shape.Dim(static_cast<size_t>(d) + 1);
}
const int64_t operand_axis_stride = operand_stride[static_cast<size_t>(axis)];
const int64_t total_updates = update_shape.NumElements();
std::vector<int32_t> base(static_cast<size_t>(total_updates));
for (int64_t f = 0; f < total_updates; ++f) {
int64_t flat = 0;
for (int64_t d = 0; d < Rank; ++d) {
if (d == axis) {
continue;
}
const int64_t coord = (f / update_stride[static_cast<size_t>(d)]) % update_shape.Dim(static_cast<size_t>(d));
flat += coord * operand_stride[static_cast<size_t>(d)];
}
base[static_cast<size_t>(f)] = static_cast<int32_t>(flat);
}
Array base_arr = MakeInt32Constant(base, update_shape);
Array stride_arr = MakeInt32Constant({static_cast<int32_t>(operand_axis_stride)}, Shape({}));
Array linear = idx * stride_arr + base_arr;
Array flat_operand = operand.Reshape(Shape({shape.NumElements()}));
Array flat_linear = linear.Reshape(Shape({total_updates}));
Array flat_values = vals.Reshape(Shape({total_updates}));
return Scatter(flat_operand, flat_linear, flat_values, 0, update).Reshape(shape);
}
TopKResult TopK(const Array& operand, int64_t k, int64_t axis, bool largest) {
const int64_t rank = static_cast<int64_t>(operand.shape().Rank());
TENZR_CHECK(rank >= 1, "plax::top_k requires rank >= 1");
if (axis < 0) {
axis += rank;
}
TENZR_CHECK(axis >= 0 && axis < rank, "plax::top_k axis out of range");
const int64_t axis_dim = operand.shape().Dim(static_cast<size_t>(axis));
TENZR_CHECK(k >= 0 && k <= axis_dim, "plax::top_k k out of range");
// Reuses the native argsort + take_along_axis primitives: indices are the first
// k of the (stable) sort permutation; values are gathered along the axis.
Array perm = Argsort(operand, axis, largest);
std::vector<Slice> slices(static_cast<size_t>(rank), Slice::All());
slices[static_cast<size_t>(axis)] = Slice(0, k, 1);
Array indices = perm.slice(slices);
Array values = TakeAlongAxis(operand, indices, axis);
return {values, indices};
}
Array Reverse(const Array& operand, const std::vector<int64_t>& dimensions) {
const int64_t rank = static_cast<int64_t>(operand.shape().Rank());
Array result = operand;
for (int64_t d : dimensions) {
const int64_t ax = d < 0 ? d + rank : d;
TENZR_CHECK(ax >= 0 && ax < rank, "plax::reverse dimension out of range");
const Shape& shape = result.shape();
const int64_t n = shape.Dim(static_cast<size_t>(ax));
// out[i] = in[n - 1 - i] == take_along_axis(in, (n - 1) - iota_axis, axis).
Array iota = MakeIotaArray(shape, ax);
Array last = MakeInt32Constant({static_cast<int32_t>(n - 1)}, Shape({}));
Array reversed_index = last - iota;
result = TakeAlongAxis(result, reversed_index, ax);
}
return result;
}
Array ConvGeneral(const Array& lhs, const Array& rhs, const std::vector<int64_t>& window_strides, const std::vector<int64_t>& lhs_dilation,
const std::vector<int64_t>& rhs_dilation, const std::vector<int64_t>& padding_low, const std::vector<int64_t>& padding_high,
const std::vector<int64_t>& lhs_spec, const std::vector<int64_t>& rhs_spec, const std::vector<int64_t>& out_spec) {
TENZR_CHECK(lhs.shape().Rank() == 4 && rhs.shape().Rank() == 4, "plax::conv_general requires rank-4 operands");
TENZR_CHECK(lhs.Dtype() == rhs.Dtype(), "plax::conv_general operand dtypes must match");
TENZR_CHECK(lhs_spec.size() == 4 && rhs_spec.size() == 4 && out_spec.size() == 4, "plax::conv_general dimension specs must have length 4");
const std::vector<int64_t> strides = window_strides.empty() ? std::vector<int64_t>{1, 1} : window_strides;
const std::vector<int64_t> ldil = lhs_dilation.empty() ? std::vector<int64_t>{1, 1} : lhs_dilation;
const std::vector<int64_t> rdil = rhs_dilation.empty() ? std::vector<int64_t>{1, 1} : rhs_dilation;
const std::vector<int64_t> plo = padding_low.empty() ? std::vector<int64_t>{0, 0} : padding_low;
const std::vector<int64_t> phi = padding_high.empty() ? std::vector<int64_t>{0, 0} : padding_high;
std::vector<int64_t> out_dims(4, 0);
for (size_t k = 0; k < 2; ++k) {
const int64_t in_d = lhs.shape().Dim(static_cast<size_t>(lhs_spec[2 + k]));
const int64_t k_d = rhs.shape().Dim(static_cast<size_t>(rhs_spec[2 + k]));
const int64_t lhs_dilated = (in_d - 1) * ldil[k] + 1;
const int64_t rhs_dilated = (k_d - 1) * rdil[k] + 1;
const int64_t numer = lhs_dilated + plo[k] + phi[k] - rhs_dilated;
TENZR_CHECK(numer >= 0, "plax::conv_general kernel is larger than the padded input");
out_dims[static_cast<size_t>(out_spec[2 + k])] = numer / strides[k] + 1;
}
out_dims[static_cast<size_t>(out_spec[0])] = lhs.shape().Dim(static_cast<size_t>(lhs_spec[0]));
out_dims[static_cast<size_t>(out_spec[1])] = rhs.shape().Dim(static_cast<size_t>(rhs_spec[0]));
Shape out_shape(out_dims);
if (txpr::TracerContext::IsActive() || lhs.IsSymbolic() || rhs.IsSymbolic()) {
auto* tape = txpr::TracerContext::CurrentTxpr();
TENZR_CHECK(tape != nullptr, "Active tape required for symbolic conv_general");
txpr::VarId l = lhs.ObtainVarId(tape);
txpr::VarId r = rhs.ObtainVarId(tape);
txpr::AttributePayload attrs;
attrs.window_strides = strides;
attrs.lhs_dilation = ldil;
attrs.rhs_dilation = rdil;
attrs.edge_padding_low = plo;
attrs.edge_padding_high = phi;
attrs.conv_lhs_spec = lhs_spec;
attrs.conv_rhs_spec = rhs_spec;
attrs.conv_out_spec = out_spec;
txpr::VarId out_id = tape->AppendInstruction(txpr::OpCode::kConvolution, {l, r}, out_shape, lhs.Dtype(), std::move(attrs));
return Array::Symbolic(out_id, out_shape, lhs.Dtype());
}
TENZR_CHECK((lhs_spec == std::vector<int64_t>{0, 1, 2, 3} && rhs_spec == std::vector<int64_t>{0, 1, 2, 3} && out_spec == std::vector<int64_t>{0, 1, 2, 3}),
"plax::conv_general eager path supports identity dimension specs only");
return Array(std::make_shared<ConvolutionNode>(lhs.node(), rhs.node(), strides, ldil, rdil, plo, phi));
}
Array Conv2d(const Array& input, const Array& weight, std::pair<int64_t, int64_t> stride, std::pair<int64_t, int64_t> padding) {
TENZR_CHECK(input.shape().Rank() == 4, "plax::conv2d input must have rank 4 [B, C, H, W]");
TENZR_CHECK(weight.shape().Rank() == 4, "plax::conv2d weight must have rank 4 [OC, C, KH, KW]");
TENZR_CHECK(input.Dtype() == weight.Dtype(), "plax::conv2d input and weight dtypes must match");
TENZR_CHECK(weight.shape().Dim(1) == input.shape().Dim(1), "plax::conv2d weight channel count must match input channels");
TENZR_CHECK(stride.first > 0 && stride.second > 0, "plax::conv2d stride must be positive");
TENZR_CHECK(padding.first >= 0 && padding.second >= 0, "plax::conv2d padding must be non-negative");
return ConvGeneral(input, weight, {stride.first, stride.second}, {1, 1}, {1, 1}, {padding.first, padding.second}, {padding.first, padding.second});
}
// Tensor Contraction & Structural Primitives
Array DotGeneral(const Array& lhs, const Array& rhs, const std::vector<int64_t>& lhs_contracting_dims, const std::vector<int64_t>& rhs_contracting_dims,
const std::vector<int64_t>& lhs_batch_dims, const std::vector<int64_t>& rhs_batch_dims) {
if (lhs.Dtype() != rhs.Dtype()) {
TENZR_THROW("plax::dot_general requires identical dtypes: " + std::string(tenzr::ToString(lhs.Dtype())) + " vs " +
std::string(tenzr::ToString(rhs.Dtype())));
}
const int64_t rank_l = static_cast<int64_t>(lhs.shape().Rank());
const int64_t rank_r = static_cast<int64_t>(rhs.shape().Rank());
auto norm = [](int64_t d, int64_t rank) { return d < 0 ? d + rank : d; };
auto norm_vec = [&](const std::vector<int64_t>& dims, int64_t rank) {
std::vector<int64_t> out(dims.size());
for (size_t i = 0; i < dims.size(); ++i) {
out[i] = norm(dims[i], rank);
TENZR_CHECK(out[i] >= 0 && out[i] < rank, "plax::dot_general dimension out of range");
}
return out;
};
std::vector<int64_t> lc = norm_vec(lhs_contracting_dims, rank_l);
std::vector<int64_t> rc = norm_vec(rhs_contracting_dims, rank_r);
std::vector<int64_t> lb = norm_vec(lhs_batch_dims, rank_l);
std::vector<int64_t> rb = norm_vec(rhs_batch_dims, rank_r);
TENZR_CHECK(lc.size() == rc.size(), "plax::dot_general contracting dim counts must match");
TENZR_CHECK(lb.size() == rb.size(), "plax::dot_general batch dim counts must match");
for (size_t i = 0; i < lb.size(); ++i) {
TENZR_CHECK(lhs.shape().Dim(static_cast<size_t>(lb[i])) == rhs.shape().Dim(static_cast<size_t>(rb[i])),
"plax::dot_general batch dimensions must have equal sizes");
}
auto contains = [](const std::vector<int64_t>& v, int64_t d) { return std::find(v.begin(), v.end(), d) != v.end(); };
std::vector<int64_t> lf;
for (int64_t d = 0; d < rank_l; ++d) {
if (!contains(lb, d) && !contains(lc, d)) {
lf.push_back(d);
}
}
std::vector<int64_t> rf;
for (int64_t d = 0; d < rank_r; ++d) {
if (!contains(rb, d) && !contains(rc, d)) {
rf.push_back(d);
}
}
std::vector<int64_t> perm_l = lb;
perm_l.insert(perm_l.end(), lf.begin(), lf.end());
perm_l.insert(perm_l.end(), lc.begin(), lc.end());
std::vector<int64_t> perm_r = rb;
perm_r.insert(perm_r.end(), rc.begin(), rc.end());
perm_r.insert(perm_r.end(), rf.begin(), rf.end());
auto is_identity = [](const std::vector<int64_t>& perm) {
for (size_t i = 0; i < perm.size(); ++i) {
if (perm[i] != static_cast<int64_t>(i)) {
return false;
}
}
return true;
};
Array lhs_p = is_identity(perm_l) ? lhs : lhs.Permute(perm_l);
Array rhs_p = is_identity(perm_r) ? rhs : rhs.Permute(perm_r);
int64_t batch = 1;
for (int64_t d : lb) {
batch *= lhs.shape().Dim(static_cast<size_t>(d));
}
int64_t m = 1;
for (int64_t d : lf) {
m *= lhs.shape().Dim(static_cast<size_t>(d));
}
int64_t k = 1;
for (int64_t d : lc) {
k *= lhs.shape().Dim(static_cast<size_t>(d));
}
int64_t n = 1;
for (int64_t d : rf) {
n *= rhs.shape().Dim(static_cast<size_t>(d));
}
Array lhs_r = lhs_p.Reshape(Shape({batch, m, k}));
Array rhs_r = rhs_p.Reshape(Shape({batch, k, n}));
Array out3 = lhs_r.Matmul(rhs_r);
std::vector<int64_t> out_dims;
out_dims.reserve(lb.size() + lf.size() + rf.size());
for (int64_t d : lb) {
out_dims.push_back(lhs.shape().Dim(static_cast<size_t>(d)));
}
for (int64_t d : lf) {
out_dims.push_back(lhs.shape().Dim(static_cast<size_t>(d)));
}
for (int64_t d : rf) {
out_dims.push_back(rhs.shape().Dim(static_cast<size_t>(d)));
}
Shape out_shape(out_dims);
if (out3.shape() == out_shape) {
return out3;
}
return out3.Reshape(out_shape);
}
Array BroadcastInDim(const Array& operand, const Shape& target_shape, const std::vector<int64_t>& broadcast_dimensions) {
(void)broadcast_dimensions;
return operand.BroadcastTo(target_shape);
}
Array Reshape(const Array& operand, const Shape& new_shape) { return operand.Reshape(new_shape); }
Array Transpose(const Array& operand, const std::vector<int64_t>& permutation) { return operand.Permute(permutation); }
Array ConvertElementType(const Array& operand, DType new_dtype) { return operand.Astype(new_dtype); }
Array Select(const Array& pred, const Array& on_true, const Array& on_false) {
if (pred.Dtype() != DType::kBool) {
TENZR_THROW("plax::select requires predicate to have DType::Bool");
}
if (on_true.Dtype() != on_false.Dtype()) {
TENZR_THROW("plax::select requires matching true/false branch dtypes");
}
if (pred.shape() != on_true.shape() || on_true.shape() != on_false.shape()) {
TENZR_THROW("plax::select requires identical shapes for predicate and branches");
}
return Array::CustomOp("select", {pred, on_true, on_false});
}
Array Clamp(const Array& min_val, const Array& operand, const Array& max_val) { return operand.Clamp(min_val, max_val); }
// Distributed Collectives & Sharding Primitives
Array AllReduce(const Array& x, std::string_view op, const std::vector<std::vector<int64_t>>& replica_groups) {
if (txpr::TracerContext::IsActive() || x.IsSymbolic()) {
auto* tape = txpr::TracerContext::CurrentTxpr();
TENZR_CHECK(tape != nullptr, "Active tape required for symbolic all_reduce");
txpr::VarId in_id = x.ObtainVarId(tape);
txpr::AttributePayload attrs;
attrs.collective_op = std::string(op);
attrs.replica_groups = replica_groups;
txpr::VarId out_id = tape->AppendInstruction(txpr::OpCode::kAllReduce, {in_id}, x.shape(), x.Dtype(), std::move(attrs));
return Array::Symbolic(out_id, x.shape(), x.Dtype());
}
// Concrete local execution semantics:
// On single replica / host, reduction is an identity copy
return x.Clone();
}
Array AllGather(const Array& x, int64_t all_gather_dim, const std::vector<std::vector<int64_t>>& replica_groups) {
int64_t group_size = replica_groups.empty() ? 1 : static_cast<int64_t>(replica_groups[0].size());
if (group_size <= 0) {
group_size = 1;
}
int64_t rank = static_cast<int64_t>(x.shape().Rank());
if (all_gather_dim < 0) {
all_gather_dim += rank;
}
TENZR_CHECK(all_gather_dim >= 0 && all_gather_dim < rank, "Invalid all_gather_dim");
auto out_dims = x.shape().Dims();
out_dims[static_cast<size_t>(all_gather_dim)] *= group_size;
Shape out_shape(std::move(out_dims));
if (txpr::TracerContext::IsActive() || x.IsSymbolic()) {
auto* tape = txpr::TracerContext::CurrentTxpr();
TENZR_CHECK(tape != nullptr, "Active tape required for symbolic all_gather");
txpr::VarId in_id = x.ObtainVarId(tape);
txpr::AttributePayload attrs;
attrs.collective_dim = all_gather_dim;
attrs.replica_groups = replica_groups;
txpr::VarId out_id = tape->AppendInstruction(txpr::OpCode::kAllGather, {in_id}, out_shape, x.Dtype(), std::move(attrs));
return Array::Symbolic(out_id, out_shape, x.Dtype());
}
if (group_size == 1) {
return x.Clone();
}
Array res = Array::Zeros(out_shape, x.Dtype());
std::vector<float> in_data = x.ToVector<float>();
std::vector<float> out_data(static_cast<size_t>(out_shape.NumElements()));
int64_t outer_size = 1;
for (int64_t i = 0; i < all_gather_dim; ++i) {
outer_size *= x.shape().Dim(static_cast<size_t>(i));
}
int64_t axis_dim = x.shape().Dim(static_cast<size_t>(all_gather_dim));
int64_t inner_size = 1;
for (size_t i = static_cast<size_t>(all_gather_dim + 1); i < static_cast<size_t>(rank); ++i) {
inner_size *= x.shape().Dim(i);
}
for (int64_t o = 0; o < outer_size; ++o) {
for (int64_t g = 0; g < group_size; ++g) {
size_t dst_offset = static_cast<size_t>((o * (axis_dim * group_size) + g * axis_dim) * inner_size);
size_t src_offset = static_cast<size_t>(o * axis_dim * inner_size);
size_t copy_elems = static_cast<size_t>(axis_dim * inner_size);
if (src_offset + copy_elems <= in_data.size() && dst_offset + copy_elems <= out_data.size()) {
std::copy_n(in_data.data() + src_offset, copy_elems, out_data.data() + dst_offset);
}
}
}
res.SetData(out_data);
return res;
}
Array AllToAll(const Array& x, int64_t split_dim, int64_t concat_dim, int64_t split_count, const std::vector<std::vector<int64_t>>& replica_groups) {
int64_t rank = static_cast<int64_t>(x.shape().Rank());
if (split_dim < 0) {
split_dim += rank;
}
if (concat_dim < 0) {
concat_dim += rank;
}
TENZR_CHECK(split_dim >= 0 && split_dim < rank, "Invalid split_dim for all_to_all");
TENZR_CHECK(concat_dim >= 0 && concat_dim < rank, "Invalid concat_dim for all_to_all");
TENZR_CHECK(split_count > 0, "split_count must be > 0");
auto out_dims = x.shape().Dims();
TENZR_CHECK(out_dims[static_cast<size_t>(split_dim)] % split_count == 0, "split_dim dimension size must be divisible by split_count");
out_dims[static_cast<size_t>(split_dim)] /= split_count;
out_dims[static_cast<size_t>(concat_dim)] *= split_count;
Shape out_shape(std::move(out_dims));
if (txpr::TracerContext::IsActive() || x.IsSymbolic()) {
auto* tape = txpr::TracerContext::CurrentTxpr();
TENZR_CHECK(tape != nullptr, "Active tape required for symbolic all_to_all");
txpr::VarId in_id = x.ObtainVarId(tape);
txpr::AttributePayload attrs;
attrs.split_dim = split_dim;
attrs.concat_dim = concat_dim;
attrs.split_count = split_count;
attrs.replica_groups = replica_groups;
txpr::VarId out_id = tape->AppendInstruction(txpr::OpCode::kAllToAll, {in_id}, out_shape, x.Dtype(), std::move(attrs));
return Array::Symbolic(out_id, out_shape, x.Dtype());
}
return x.Reshape(out_shape);
}
Array ReduceScatter(const Array& x, std::string_view op, int64_t scatter_dim, const std::vector<std::vector<int64_t>>& replica_groups) {
int64_t group_size = replica_groups.empty() ? 1 : static_cast<int64_t>(replica_groups[0].size());
if (group_size <= 0) {
group_size = 1;
}
int64_t rank = static_cast<int64_t>(x.shape().Rank());
if (scatter_dim < 0) {
scatter_dim += rank;
}
TENZR_CHECK(scatter_dim >= 0 && scatter_dim < rank, "Invalid scatter_dim for reduce_scatter");
auto out_dims = x.shape().Dims();
TENZR_CHECK(out_dims[static_cast<size_t>(scatter_dim)] % group_size == 0, "scatter_dim dimension size must be divisible by group_size");
out_dims[static_cast<size_t>(scatter_dim)] /= group_size;
Shape out_shape(std::move(out_dims));
if (txpr::TracerContext::IsActive() || x.IsSymbolic()) {
auto* tape = txpr::TracerContext::CurrentTxpr();
TENZR_CHECK(tape != nullptr, "Active tape required for symbolic reduce_scatter");
txpr::VarId in_id = x.ObtainVarId(tape);
txpr::AttributePayload attrs;
attrs.collective_op = std::string(op);
attrs.collective_dim = scatter_dim;
attrs.replica_groups = replica_groups;
txpr::VarId out_id = tape->AppendInstruction(txpr::OpCode::kReduceScatter, {in_id}, out_shape, x.Dtype(), std::move(attrs));
return Array::Symbolic(out_id, out_shape, x.Dtype());
}
std::vector<Slice> slices(static_cast<size_t>(rank));
slices[static_cast<size_t>(scatter_dim)] = Slice(0, out_shape[static_cast<size_t>(scatter_dim)]);
return x.slice(slices);
}
Array Psum(const Array& x, std::string_view /*axis_name*/) { return AllReduce(x, "sum"); }
Array Pmean(const Array& x, std::string_view /*axis_name*/) { return AllReduce(x, "sum"); }
Array WithShardingConstraint(const Array& x, const NamedSharding& sharding) {
Array res = x;
res.Shard(sharding);
return res;
}
} // namespace tenzr::plax