// 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/txpr/vmap.h"
#include <algorithm>
#include <unordered_map>
#include <vector>
#include "tenzr/common/error.h"
#include "tenzr/common/trace.h"
#include "tenzr/txpr/txpr.h"
namespace tenzr::txpr {
Txpr BatchTxpr(const Txpr& orig, int64_t batch_size, const std::vector<int>& in_axes, int out_axis) {
TENZR_TRACE_SCOPE_CAT("vmap", "txpr::batch_txpr");
TENZR_CHECK(batch_size > 0, "batch_txpr requires positive batch size");
(void)out_axis; // Default batch output at axis 0
Txpr batched;
std::unordered_map<uint32_t, VarId> var_map;
std::unordered_map<uint32_t, int> bdim_map; // maps orig VarId to its batch dimension in batched Txpr (-1 if unbatched)
// 1. Map inputs
const auto& in_binders = orig.InBinders();
for (size_t i = 0; i < in_binders.size(); ++i) {
VarId orig_vid = in_binders[i];
const auto& orig_var = orig.Variable(orig_vid);
int ax = (i < in_axes.size()) ? in_axes[i] : 0;
if (ax >= 0) {
// Batched input: prepend batch dimension at axis 0
std::vector<int64_t> new_dims;
new_dims.reserve(orig_var.shape.Rank() + 1);
new_dims.push_back(batch_size);
for (size_t d = 0; d < orig_var.shape.Rank(); ++d) {
new_dims.push_back(orig_var.shape.Dim(d));
}
VarId batched_vid = batched.AddVar(Shape(new_dims), orig_var.dtype, orig_var.name);
batched.MarkInput(batched_vid);
var_map[orig_vid.id] = batched_vid;
bdim_map[orig_vid.id] = 0;
} else {
// Unbatched input
VarId batched_vid = batched.AddVar(orig_var.shape, orig_var.dtype, orig_var.name);
batched.MarkInput(batched_vid);
var_map[orig_vid.id] = batched_vid;
bdim_map[orig_vid.id] = -1;
}
}
// 2. Transform equations
for (const auto& eqn : orig.Equations()) {
std::vector<VarId> mapped_ins;
mapped_ins.reserve(eqn.in_vars.size());
for (const auto& in_v : eqn.in_vars) {
mapped_ins.push_back(var_map.at(in_v.id));
}
const auto& orig_out = orig.Variable(eqn.out_vars[0]);
switch (eqn.op) {
// Activations & Type Conversion & elementwise math (via CustomOp)
case OpCode::kRelu:
case OpCode::kSigmoid:
case OpCode::kConvert:
case OpCode::kCustomOp: {
int in_bdim = eqn.in_vars.empty() ? -1 : bdim_map.at(eqn.in_vars[0].id);
if (in_bdim < 0) {
VarId out_v = batched.AppendInstruction(eqn.op, mapped_ins, orig_out.shape, orig_out.dtype, eqn.attrs);
var_map[eqn.out_vars[0].id] = out_v;
bdim_map[eqn.out_vars[0].id] = -1;
} else {
std::vector<int64_t> out_dims;
out_dims.reserve(orig_out.shape.Rank() + 1);
out_dims.push_back(batch_size);
for (size_t d = 0; d < orig_out.shape.Rank(); ++d) {
out_dims.push_back(orig_out.shape.Dim(d));
}
VarId out_v = batched.AppendInstruction(eqn.op, mapped_ins, Shape(out_dims), orig_out.dtype, eqn.attrs);
var_map[eqn.out_vars[0].id] = out_v;
bdim_map[eqn.out_vars[0].id] = 0;
}
break;
}
// Binary Elementwise
case OpCode::kAdd:
case OpCode::kSub:
case OpCode::kMul:
case OpCode::kDiv: {
bool any_batched = false;
for (const auto& in_v : eqn.in_vars) {
if (bdim_map.at(in_v.id) >= 0) {
any_batched = true;
break;
}
}
if (!any_batched) {
VarId out_v = batched.AppendInstruction(eqn.op, mapped_ins, orig_out.shape, orig_out.dtype, eqn.attrs);
var_map[eqn.out_vars[0].id] = out_v;
bdim_map[eqn.out_vars[0].id] = -1;
} else {
std::vector<int64_t> out_dims;
out_dims.reserve(orig_out.shape.Rank() + 1);
out_dims.push_back(batch_size);
for (size_t d = 0; d < orig_out.shape.Rank(); ++d) {
out_dims.push_back(orig_out.shape.Dim(d));
}
VarId out_v = batched.AppendInstruction(eqn.op, mapped_ins, Shape(out_dims), orig_out.dtype, eqn.attrs);
var_map[eqn.out_vars[0].id] = out_v;
bdim_map[eqn.out_vars[0].id] = 0;
}
break;
}
// Reductions
case OpCode::kReduceSum:
case OpCode::kReduceProd:
case OpCode::kReduceMax:
case OpCode::kReduceMin: {
int in_bdim = bdim_map.at(eqn.in_vars[0].id);
if (in_bdim < 0) {
VarId out_v = batched.AppendInstruction(eqn.op, mapped_ins, orig_out.shape, orig_out.dtype, eqn.attrs);
var_map[eqn.out_vars[0].id] = out_v;
bdim_map[eqn.out_vars[0].id] = -1;
} else {
AttributePayload batched_attrs = eqn.attrs;
const auto& in_orig_var = orig.Variable(eqn.in_vars[0]);
size_t orig_rank = in_orig_var.shape.Rank();
std::vector<int64_t> shifted_axes;
if (eqn.attrs.axes.empty()) {
// Reducing all dimensions in original tensor:
// In batched tensor, shift all axes by +1 (axes 1 .. orig_rank)
shifted_axes.reserve(orig_rank);
for (size_t a = 1; a <= orig_rank; ++a) {
shifted_axes.push_back(static_cast<int64_t>(a));
}
} else {
// Specific axes reduced: shift each axis by +1
shifted_axes.reserve(eqn.attrs.axes.size());
for (int64_t ax : eqn.attrs.axes) {
shifted_axes.push_back(ax + 1);
}
}
batched_attrs.axes = shifted_axes;
// Compute batched output shape
std::vector<int64_t> out_dims;
out_dims.push_back(batch_size);
if (eqn.attrs.keepdims) {
for (size_t d = 0; d < orig_out.shape.Rank(); ++d) {
out_dims.push_back(orig_out.shape.Dim(d));
}
} else {
if (orig_out.shape.Rank() > 0) {
for (size_t d = 0; d < orig_out.shape.Rank(); ++d) {
out_dims.push_back(orig_out.shape.Dim(d));
}
}
}
VarId out_v = batched.AppendInstruction(eqn.op, mapped_ins, Shape(out_dims), orig_out.dtype, batched_attrs);
var_map[eqn.out_vars[0].id] = out_v;
bdim_map[eqn.out_vars[0].id] = 0;
}
break;
}
case OpCode::kReduceWindow: {
int in_bdim = bdim_map.at(eqn.in_vars[0].id);
if (in_bdim < 0) {
VarId out_v = batched.AppendInstruction(eqn.op, mapped_ins, orig_out.shape, orig_out.dtype, eqn.attrs);
var_map[eqn.out_vars[0].id] = out_v;
bdim_map[eqn.out_vars[0].id] = -1;
break;
}
// The window/padding specs are per operand dimension; prepend the batch.
AttributePayload batched_attrs = eqn.attrs;
batched_attrs.window_dimensions.insert(batched_attrs.window_dimensions.begin(), 1);
batched_attrs.window_strides.insert(batched_attrs.window_strides.begin(), 1);
batched_attrs.base_dilations.insert(batched_attrs.base_dilations.begin(), 1);
batched_attrs.window_dilations.insert(batched_attrs.window_dilations.begin(), 1);
batched_attrs.edge_padding_low.insert(batched_attrs.edge_padding_low.begin(), 0);
batched_attrs.edge_padding_high.insert(batched_attrs.edge_padding_high.begin(), 0);
std::vector<int64_t> out_dims;
out_dims.reserve(orig_out.shape.Rank() + 1);
out_dims.push_back(batch_size);
for (size_t d = 0; d < orig_out.shape.Rank(); ++d) {
out_dims.push_back(orig_out.shape.Dim(d));
}
VarId out_v = batched.AppendInstruction(eqn.op, mapped_ins, Shape(out_dims), orig_out.dtype, batched_attrs);
var_map[eqn.out_vars[0].id] = out_v;
bdim_map[eqn.out_vars[0].id] = 0;
break;
}
case OpCode::kConvolution: {
int lhs_bdim = bdim_map.at(eqn.in_vars[0].id);
int rhs_bdim = bdim_map.at(eqn.in_vars[1].id);
if (lhs_bdim < 0 && rhs_bdim < 0) {
VarId out_v = batched.AppendInstruction(eqn.op, mapped_ins, orig_out.shape, orig_out.dtype, eqn.attrs);
var_map[eqn.out_vars[0].id] = out_v;
bdim_map[eqn.out_vars[0].id] = -1;
break;
}
TENZR_CHECK(rhs_bdim < 0, "vmap over convolution with a batched kernel is not supported");
// Fold the batch into the operand's N dim (the kernel is shared):
// [B,N,C,H,W] -> [B*N,C,H,W] -> conv -> [B,N,O,OH,OW].
const auto& lhs_var = orig.Variable(eqn.in_vars[0]);
const int64_t n = lhs_var.shape.Dim(0);
std::vector<int64_t> flat_dims = {batch_size * n, lhs_var.shape.Dim(1), lhs_var.shape.Dim(2), lhs_var.shape.Dim(3)};
VarId lhs_flat = batched.AppendInstruction(OpCode::kReshape, {mapped_ins[0]}, Shape(flat_dims), lhs_var.dtype, {});
std::vector<int64_t> conv_dims = {batch_size * n, orig_out.shape.Dim(1), orig_out.shape.Dim(2), orig_out.shape.Dim(3)};
VarId conv_id = batched.AppendInstruction(eqn.op, {lhs_flat, mapped_ins[1]}, Shape(conv_dims), orig_out.dtype, eqn.attrs);
std::vector<int64_t> out_dims;
out_dims.reserve(orig_out.shape.Rank() + 1);
out_dims.push_back(batch_size);
for (size_t d = 0; d < orig_out.shape.Rank(); ++d) {
out_dims.push_back(orig_out.shape.Dim(d));
}
VarId out_v = batched.AppendInstruction(OpCode::kReshape, {conv_id}, Shape(out_dims), orig_out.dtype, {});
var_map[eqn.out_vars[0].id] = out_v;
bdim_map[eqn.out_vars[0].id] = 0;
break;
}
case OpCode::kDynamicSlice:
case OpCode::kDynamicUpdateSlice: {
TENZR_THROW("vmap over dynamic_slice/dynamic_update_slice is not supported");
}
// Matrix Multiplication
case OpCode::kMatmul: {
int lhs_bdim = bdim_map.at(eqn.in_vars[0].id);
int rhs_bdim = bdim_map.at(eqn.in_vars[1].id);
VarId in_lhs = mapped_ins[0];
VarId in_rhs = mapped_ins[1];
if (lhs_bdim >= 0 && rhs_bdim < 0) {
// rhs is unbatched, broadcast along new leading axis 0
const auto& rhs_var = orig.Variable(eqn.in_vars[1]);
std::vector<int64_t> bcast_dims = {batch_size};
for (size_t d = 0; d < rhs_var.shape.Rank(); ++d) {
bcast_dims.push_back(rhs_var.shape.Dim(d));
}
AttributePayload bcast_attrs;
std::vector<int64_t> bcast_axes;
for (size_t d = 0; d < rhs_var.shape.Rank(); ++d) {
bcast_axes.push_back(static_cast<int64_t>(d + 1));
}
bcast_attrs.axes = bcast_axes;
in_rhs = batched.AppendInstruction(OpCode::kBroadcastTo, {in_rhs}, Shape(bcast_dims), rhs_var.dtype, bcast_attrs);
} else if (lhs_bdim < 0 && rhs_bdim >= 0) {
// lhs is unbatched, broadcast along new leading axis 0
const auto& lhs_var = orig.Variable(eqn.in_vars[0]);
std::vector<int64_t> bcast_dims = {batch_size};
for (size_t d = 0; d < lhs_var.shape.Rank(); ++d) {
bcast_dims.push_back(lhs_var.shape.Dim(d));
}
AttributePayload bcast_attrs;
std::vector<int64_t> bcast_axes;
for (size_t d = 0; d < lhs_var.shape.Rank(); ++d) {
bcast_axes.push_back(static_cast<int64_t>(d + 1));
}
bcast_attrs.axes = bcast_axes;
in_lhs = batched.AppendInstruction(OpCode::kBroadcastTo, {in_lhs}, Shape(bcast_dims), lhs_var.dtype, bcast_attrs);
}
if (lhs_bdim < 0 && rhs_bdim < 0) {
VarId out_v = batched.AppendInstruction(eqn.op, {in_lhs, in_rhs}, orig_out.shape, orig_out.dtype, eqn.attrs);
var_map[eqn.out_vars[0].id] = out_v;
bdim_map[eqn.out_vars[0].id] = -1;
} else {
std::vector<int64_t> out_dims;
out_dims.reserve(orig_out.shape.Rank() + 1);
out_dims.push_back(batch_size);
for (size_t d = 0; d < orig_out.shape.Rank(); ++d) {
out_dims.push_back(orig_out.shape.Dim(d));
}
VarId out_v = batched.AppendInstruction(eqn.op, {in_lhs, in_rhs}, Shape(out_dims), orig_out.dtype, eqn.attrs);
var_map[eqn.out_vars[0].id] = out_v;
bdim_map[eqn.out_vars[0].id] = 0;
}
break;
}
// Constants
case OpCode::kConstant: {
VarId out_v = batched.AppendInstruction(eqn.op, {}, orig_out.shape, orig_out.dtype, eqn.attrs);
var_map[eqn.out_vars[0].id] = out_v;
bdim_map[eqn.out_vars[0].id] = -1;
break;
}
// Shape & View transforms
case OpCode::kReshape: {
int in_bdim = bdim_map.at(eqn.in_vars[0].id);
if (in_bdim < 0) {
VarId out_v = batched.AppendInstruction(eqn.op, mapped_ins, orig_out.shape, orig_out.dtype, eqn.attrs);
var_map[eqn.out_vars[0].id] = out_v;
bdim_map[eqn.out_vars[0].id] = -1;
} else {
std::vector<int64_t> out_dims;
out_dims.reserve(orig_out.shape.Rank() + 1);
out_dims.push_back(batch_size);
for (size_t d = 0; d < orig_out.shape.Rank(); ++d) {
out_dims.push_back(orig_out.shape.Dim(d));
}
VarId out_v = batched.AppendInstruction(eqn.op, mapped_ins, Shape(out_dims), orig_out.dtype, eqn.attrs);
var_map[eqn.out_vars[0].id] = out_v;
bdim_map[eqn.out_vars[0].id] = 0;
}
break;
}
case OpCode::kTranspose: {
int in_bdim = bdim_map.at(eqn.in_vars[0].id);
if (in_bdim < 0) {
VarId out_v = batched.AppendInstruction(eqn.op, mapped_ins, orig_out.shape, orig_out.dtype, eqn.attrs);
var_map[eqn.out_vars[0].id] = out_v;
bdim_map[eqn.out_vars[0].id] = -1;
} else {
AttributePayload batched_attrs = eqn.attrs;
std::vector<int64_t> batched_perm;
batched_perm.reserve(eqn.attrs.permutation.size() + 1);
batched_perm.push_back(0); // Batch axis remains at 0
for (int64_t p : eqn.attrs.permutation) {
batched_perm.push_back(p + 1);
}
batched_attrs.permutation = batched_perm;
std::vector<int64_t> out_dims;
out_dims.reserve(orig_out.shape.Rank() + 1);
out_dims.push_back(batch_size);
for (size_t d = 0; d < orig_out.shape.Rank(); ++d) {
out_dims.push_back(orig_out.shape.Dim(d));
}
VarId out_v = batched.AppendInstruction(eqn.op, mapped_ins, Shape(out_dims), orig_out.dtype, batched_attrs);
var_map[eqn.out_vars[0].id] = out_v;
bdim_map[eqn.out_vars[0].id] = 0;
}
break;
}
case OpCode::kBroadcastTo: {
AttributePayload batched_attrs = eqn.attrs;
int in_bdim = bdim_map.at(eqn.in_vars[0].id);
const auto& in_orig_var = orig.Variable(eqn.in_vars[0]);
std::vector<int64_t> shifted_axes;
if (in_bdim >= 0) {
shifted_axes.push_back(0);
for (int64_t ax : eqn.attrs.axes) {
shifted_axes.push_back(ax + 1);
}
} else {
if (in_orig_var.shape.Rank() == 0) {
// Scalar broadcast: broadcast dimensions is empty
shifted_axes = {};
} else {
for (int64_t ax : eqn.attrs.axes) {
shifted_axes.push_back(ax + 1);
}
}
}
batched_attrs.axes = shifted_axes;
std::vector<int64_t> out_dims;
out_dims.reserve(orig_out.shape.Rank() + 1);
out_dims.push_back(batch_size);
for (size_t d = 0; d < orig_out.shape.Rank(); ++d) {
out_dims.push_back(orig_out.shape.Dim(d));
}
VarId out_v = batched.AppendInstruction(eqn.op, mapped_ins, Shape(out_dims), orig_out.dtype, batched_attrs);
var_map[eqn.out_vars[0].id] = out_v;
bdim_map[eqn.out_vars[0].id] = 0;
break;
}
case OpCode::kPad: {
int in_bdim = bdim_map.at(eqn.in_vars[0].id);
int pv_bdim = bdim_map.at(eqn.in_vars[1].id);
if (pv_bdim >= 0) {
TENZR_THROW("vmap over plax::pad with a batched padding value is not supported");
}
if (in_bdim < 0) {
VarId out_v = batched.AppendInstruction(eqn.op, mapped_ins, orig_out.shape, orig_out.dtype, eqn.attrs);
var_map[eqn.out_vars[0].id] = out_v;
bdim_map[eqn.out_vars[0].id] = -1;
break;
}
AttributePayload batched_attrs = eqn.attrs;
std::vector<int64_t> low{0};
std::vector<int64_t> high{0};
std::vector<int64_t> inter{0};
low.insert(low.end(), eqn.attrs.edge_padding_low.begin(), eqn.attrs.edge_padding_low.end());
high.insert(high.end(), eqn.attrs.edge_padding_high.begin(), eqn.attrs.edge_padding_high.end());
inter.insert(inter.end(), eqn.attrs.interior_padding.begin(), eqn.attrs.interior_padding.end());
batched_attrs.edge_padding_low = std::move(low);
batched_attrs.edge_padding_high = std::move(high);
batched_attrs.interior_padding = std::move(inter);
std::vector<int64_t> out_dims;
out_dims.reserve(orig_out.shape.Rank() + 1);
out_dims.push_back(batch_size);
for (size_t d = 0; d < orig_out.shape.Rank(); ++d) {
out_dims.push_back(orig_out.shape.Dim(d));
}
VarId out_v = batched.AppendInstruction(eqn.op, mapped_ins, Shape(out_dims), orig_out.dtype, batched_attrs);
var_map[eqn.out_vars[0].id] = out_v;
bdim_map[eqn.out_vars[0].id] = 0;
break;
}
case OpCode::kArgSort: {
int in_bdim = bdim_map.at(eqn.in_vars[0].id);
if (in_bdim < 0) {
VarId out_v = batched.AppendInstruction(eqn.op, mapped_ins, orig_out.shape, orig_out.dtype, eqn.attrs);
var_map[eqn.out_vars[0].id] = out_v;
bdim_map[eqn.out_vars[0].id] = -1;
break;
}
std::vector<int64_t> out_dims;
out_dims.reserve(orig_out.shape.Rank() + 1);
out_dims.push_back(batch_size);
for (size_t d = 0; d < orig_out.shape.Rank(); ++d) {
out_dims.push_back(orig_out.shape.Dim(d));
}
Shape batched_shape(out_dims);
const int64_t baxis = eqn.attrs.sort_dimension + 1;
AttributePayload batched_attrs = eqn.attrs;
batched_attrs.sort_dimension = baxis;
if (eqn.attrs.sort_return_values) {
VarId out_v = batched.AppendInstruction(eqn.op, {mapped_ins[0]}, batched_shape, orig_out.dtype, batched_attrs);
var_map[eqn.out_vars[0].id] = out_v;
bdim_map[eqn.out_vars[0].id] = 0;
break;
}
int64_t inner = 1;
for (size_t d = static_cast<size_t>(baxis) + 1; d < batched_shape.Rank(); ++d) {
inner *= batched_shape.Dim(d);
}
const int64_t dim = batched_shape.Dim(static_cast<size_t>(baxis));
const int64_t total = batched_shape.NumElements();
const int64_t outer = total / (dim * inner);
std::vector<int64_t> iota(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) {
iota[static_cast<size_t>((o * dim + i) * inner + in)] = i;
}
}
}
AttributePayload iota_attrs;
iota_attrs.int_tensor_data = std::move(iota);
VarId iota_id = batched.AppendInstruction(OpCode::kConstant, {}, batched_shape, DType::kInt32, std::move(iota_attrs));
VarId out_v = batched.AppendInstruction(eqn.op, {mapped_ins[0], iota_id}, batched_shape, orig_out.dtype, batched_attrs);
var_map[eqn.out_vars[0].id] = out_v;
bdim_map[eqn.out_vars[0].id] = 0;
break;
}
case OpCode::kGather: {
int in_bdim = bdim_map.at(eqn.in_vars[0].id);
int idx_bdim = bdim_map.at(eqn.in_vars[1].id);
if (idx_bdim >= 0) {
// Per-batch gather (used by plax::take_along_axis): operand and
// index share the leading batch dimension. Flatten the batch into a
// single 1-D gather by linearizing the indices with a batch offset
// b * operand_size.
TENZR_CHECK(in_bdim == idx_bdim, "vmap over gather requires the operand and indices to be batched on the same axis");
const auto& op_var = orig.Variable(eqn.in_vars[0]);
const auto& idx_var = orig.Variable(eqn.in_vars[1]);
TENZR_CHECK(op_var.shape.Rank() == 1 && eqn.attrs.gather_axis == 0,
"vmap over gather with a batched index requires a rank-1 operand gathered on axis 0");
const int64_t operand_size = op_var.shape.Dim(0);
const int64_t index_count = orig_out.shape.Dim(0);
VarId op_flat = batched.AppendInstruction(OpCode::kReshape, {mapped_ins[0]}, Shape({batch_size * operand_size}), op_var.dtype, {});
VarId idx_2d = batched.AppendInstruction(OpCode::kReshape, {mapped_ins[1]}, Shape({batch_size, index_count}), idx_var.dtype, {});
std::vector<int32_t> batch_offsets(static_cast<size_t>(batch_size * index_count));
for (int64_t b = 0; b < batch_size; ++b) {
for (int64_t i = 0; i < index_count; ++i) {
batch_offsets[static_cast<size_t>(b * index_count + i)] = static_cast<int32_t>(b * operand_size);
}
}
AttributePayload offset_attrs;
offset_attrs.int_tensor_data.assign(batch_offsets.begin(), batch_offsets.end());
VarId offset_id = batched.AppendInstruction(OpCode::kConstant, {}, Shape({batch_size, index_count}), DType::kInt32, std::move(offset_attrs));
VarId linear = batched.AppendInstruction(OpCode::kAdd, {idx_2d, offset_id}, Shape({batch_size, index_count}), DType::kInt32, {});
VarId linear_col = batched.AppendInstruction(OpCode::kReshape, {linear}, Shape({batch_size * index_count, 1}), DType::kInt32, {});
AttributePayload gather_attrs;
gather_attrs.gather_axis = 0;
VarId gathered =
batched.AppendInstruction(OpCode::kGather, {op_flat, linear_col}, Shape({batch_size * index_count}), op_var.dtype, std::move(gather_attrs));
VarId out_v = batched.AppendInstruction(OpCode::kReshape, {gathered}, Shape({batch_size, index_count}), op_var.dtype, {});
var_map[eqn.out_vars[0].id] = out_v;
bdim_map[eqn.out_vars[0].id] = 0;
break;
}
if (in_bdim < 0) {
VarId out_v = batched.AppendInstruction(eqn.op, mapped_ins, orig_out.shape, orig_out.dtype, eqn.attrs);
var_map[eqn.out_vars[0].id] = out_v;
bdim_map[eqn.out_vars[0].id] = -1;
break;
}
// stablehlo.gather emits [k, batch, ...rest]; transpose the batch to dim 0.
std::vector<int64_t> gather_dims;
gather_dims.reserve(orig_out.shape.Rank() + 1);
gather_dims.push_back(orig_out.shape.Dim(0));
gather_dims.push_back(batch_size);
for (size_t d = 1; d < orig_out.shape.Rank(); ++d) {
gather_dims.push_back(orig_out.shape.Dim(d));
}
AttributePayload batched_attrs = eqn.attrs;
batched_attrs.gather_axis = eqn.attrs.gather_axis + 1;
VarId gathered = batched.AppendInstruction(eqn.op, mapped_ins, Shape(gather_dims), orig_out.dtype, batched_attrs);
std::vector<int64_t> out_dims;
out_dims.reserve(orig_out.shape.Rank() + 1);
out_dims.push_back(batch_size);
for (size_t d = 0; d < orig_out.shape.Rank(); ++d) {
out_dims.push_back(orig_out.shape.Dim(d));
}
std::vector<int64_t> perm;
perm.reserve(gather_dims.size());
perm.push_back(1);
perm.push_back(0);
for (size_t d = 2; d < gather_dims.size(); ++d) {
perm.push_back(static_cast<int64_t>(d));
}
AttributePayload tr_attrs;
tr_attrs.permutation = std::move(perm);
VarId out_v = batched.AppendInstruction(OpCode::kTranspose, {gathered}, Shape(out_dims), orig_out.dtype, std::move(tr_attrs));
var_map[eqn.out_vars[0].id] = out_v;
bdim_map[eqn.out_vars[0].id] = 0;
break;
}
case OpCode::kScatter: {
int op_bdim = bdim_map.at(eqn.in_vars[0].id);
int idx_bdim = bdim_map.at(eqn.in_vars[1].id);
int upd_bdim = bdim_map.at(eqn.in_vars[2].id);
if (idx_bdim >= 0) {
// Per-batch scatter (used by plax::put_along_axis): operand, indices
// and updates share the leading batch dim. Linearize into a 1-D
// scatter with a b * operand_size offset.
TENZR_CHECK(op_bdim == idx_bdim, "vmap over scatter requires operand and indices batched on the same axis");
const auto& op_var = orig.Variable(eqn.in_vars[0]);
const auto& idx_var = orig.Variable(eqn.in_vars[1]);
const auto& upd_var = orig.Variable(eqn.in_vars[2]);
TENZR_CHECK(op_var.shape.Rank() == 1 && eqn.attrs.scatter_axis == 0,
"vmap over scatter with a batched index requires a rank-1 operand scattered on axis 0");
const int64_t operand_size = op_var.shape.Dim(0);
const int64_t update_count = idx_var.shape.Dim(0);
VarId op_flat = batched.AppendInstruction(OpCode::kReshape, {mapped_ins[0]}, Shape({batch_size * operand_size}), op_var.dtype, {});
VarId idx_2d = batched.AppendInstruction(OpCode::kReshape, {mapped_ins[1]}, Shape({batch_size, update_count}), idx_var.dtype, {});
VarId upd_flat;
if (upd_bdim >= 0) {
upd_flat = batched.AppendInstruction(OpCode::kReshape, {mapped_ins[2]}, Shape({batch_size * update_count}), upd_var.dtype, {});
} else {
AttributePayload b_attrs;
std::vector<int64_t> axes;
for (size_t d = 0; d < upd_var.shape.Rank(); ++d) {
axes.push_back(static_cast<int64_t>(d) + 1);
}
b_attrs.axes = std::move(axes);
VarId up_b =
batched.AppendInstruction(OpCode::kBroadcastTo, {mapped_ins[2]}, Shape({batch_size, update_count}), upd_var.dtype, std::move(b_attrs));
upd_flat = batched.AppendInstruction(OpCode::kReshape, {up_b}, Shape({batch_size * update_count}), upd_var.dtype, {});
}
std::vector<int32_t> batch_offsets(static_cast<size_t>(batch_size * update_count));
for (int64_t b = 0; b < batch_size; ++b) {
for (int64_t i = 0; i < update_count; ++i) {
batch_offsets[static_cast<size_t>(b * update_count + i)] = static_cast<int32_t>(b * operand_size);
}
}
AttributePayload offset_attrs;
offset_attrs.int_tensor_data.assign(batch_offsets.begin(), batch_offsets.end());
VarId offset_id = batched.AppendInstruction(OpCode::kConstant, {}, Shape({batch_size, update_count}), DType::kInt32, std::move(offset_attrs));
VarId linear = batched.AppendInstruction(OpCode::kAdd, {idx_2d, offset_id}, Shape({batch_size, update_count}), DType::kInt32, {});
VarId linear_col = batched.AppendInstruction(OpCode::kReshape, {linear}, Shape({batch_size * update_count, 1}), DType::kInt32, {});
AttributePayload scatter_attrs;
scatter_attrs.scatter_axis = 0;
scatter_attrs.scatter_update = eqn.attrs.scatter_update;
VarId scattered = batched.AppendInstruction(OpCode::kScatter, {op_flat, linear_col, upd_flat}, Shape({batch_size * operand_size}), op_var.dtype,
std::move(scatter_attrs));
VarId out_v = batched.AppendInstruction(OpCode::kReshape, {scattered}, Shape({batch_size, operand_size}), op_var.dtype, {});
var_map[eqn.out_vars[0].id] = out_v;
bdim_map[eqn.out_vars[0].id] = 0;
break;
}
if (op_bdim < 0) {
VarId out_v = batched.AppendInstruction(eqn.op, mapped_ins, orig_out.shape, orig_out.dtype, eqn.attrs);
var_map[eqn.out_vars[0].id] = out_v;
bdim_map[eqn.out_vars[0].id] = -1;
break;
}
TENZR_CHECK(op_bdim >= 0, "vmap over plax::scatter requires the operand to be batched");
std::vector<VarId> ins = mapped_ins;
if (upd_bdim < 0) {
const auto& upd_var = orig.Variable(eqn.in_vars[2]);
std::vector<int64_t> bd;
bd.push_back(batch_size);
for (size_t d = 0; d < upd_var.shape.Rank(); ++d) {
bd.push_back(upd_var.shape.Dim(d));
}
AttributePayload b_attrs;
std::vector<int64_t> axes;
for (size_t d = 0; d < upd_var.shape.Rank(); ++d) {
axes.push_back(static_cast<int64_t>(d) + 1);
}
b_attrs.axes = std::move(axes);
ins[2] = batched.AppendInstruction(OpCode::kBroadcastTo, {mapped_ins[2]}, Shape(bd), upd_var.dtype, std::move(b_attrs));
}
AttributePayload batched_attrs = eqn.attrs;
batched_attrs.scatter_axis = eqn.attrs.scatter_axis + 1;
std::vector<int64_t> out_dims;
out_dims.reserve(orig_out.shape.Rank() + 1);
out_dims.push_back(batch_size);
for (size_t d = 0; d < orig_out.shape.Rank(); ++d) {
out_dims.push_back(orig_out.shape.Dim(d));
}
VarId out_v = batched.AppendInstruction(eqn.op, ins, Shape(out_dims), orig_out.dtype, batched_attrs);
var_map[eqn.out_vars[0].id] = out_v;
bdim_map[eqn.out_vars[0].id] = 0;
break;
}
case OpCode::kSlice: {
AttributePayload batched_attrs = eqn.attrs;
std::vector<Slice> shifted_slices;
shifted_slices.push_back(Slice(0, batch_size, 1));
for (const auto& s : eqn.attrs.slices) {
shifted_slices.push_back(s);
}
batched_attrs.slices = shifted_slices;
std::vector<int64_t> out_dims;
out_dims.reserve(orig_out.shape.Rank() + 1);
out_dims.push_back(batch_size);
for (size_t d = 0; d < orig_out.shape.Rank(); ++d) {
out_dims.push_back(orig_out.shape.Dim(d));
}
VarId out_v = batched.AppendInstruction(eqn.op, mapped_ins, Shape(out_dims), orig_out.dtype, batched_attrs);
var_map[eqn.out_vars[0].id] = out_v;
bdim_map[eqn.out_vars[0].id] = 0;
break;
}
default: {
TENZR_THROW("Unsupported OpCode in batch_txpr: " + std::string(OpCodeToString(eqn.op)));
}
}
}
// 3. Map output binders
for (VarId out_id : orig.OutBinders()) {
batched.MarkOutput(var_map.at(out_id.id));
}
return batched;
}
} // namespace tenzr::txpr