// 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/grad.h"
#include <algorithm>
#include <numeric>
#include <unordered_map>
#include <vector>
#include "tenzr/common/error.h"
#include "tenzr/common/trace.h"
#include "tenzr/execution/execution_engine.h"
#include "tenzr/frontend/node.h"
#include "tenzr/frontend/node_traversal.h"
#include "tenzr/plax/plax.h"
#include "tenzr/txpr/custom_op_deriv.h"
#include "tenzr/txpr/custom_vjp.h"
#include "tenzr/txpr/jit.h"
#include "tenzr/txpr/node_to_txpr.h"
#include "tenzr/txpr/tracer.h"
#include "tenzr/txpr/txpr.h"
namespace tenzr::txpr {
std::unordered_map<uint64_t, Array> VjpSweep(const Array& loss, const std::vector<Array>& targets, const Array& cotangent);
namespace {
bool IsZeroGradientOpcode(txpr::OpCode op) { return op == txpr::OpCode::kConstant; }
Array unbroadcast(const Array& grad_y, const Shape& target_shape) {
if (grad_y.shape() == target_shape) {
return grad_y;
}
if (target_shape.NumElements() == 1) {
return grad_y.Sum(std::nullopt, /*keepdims=*/false).Reshape(target_shape);
}
size_t y_rank = grad_y.shape().Rank();
size_t t_rank = target_shape.Rank();
TENZR_CHECK(t_rank <= y_rank, "Target shape rank cannot exceed gradient rank in unbroadcast");
std::vector<int64_t> axes_to_reduce_keepdims;
std::vector<int64_t> axes_to_reduce_drop;
for (size_t i = 0; i < y_rank - t_rank; ++i) {
axes_to_reduce_drop.push_back(static_cast<int64_t>(i));
}
for (size_t i = 0; i < t_rank; ++i) {
size_t y_idx = (y_rank - t_rank) + i;
int64_t t_dim = target_shape.Dim(i);
int64_t y_dim = grad_y.shape().Dim(y_idx);
if (t_dim == 1 && y_dim > 1) {
axes_to_reduce_keepdims.push_back(static_cast<int64_t>(y_idx));
}
}
Array res = grad_y;
if (!axes_to_reduce_keepdims.empty()) {
res = res.Sum(axes_to_reduce_keepdims, /*keepdims=*/true);
}
if (!axes_to_reduce_drop.empty()) {
res = res.Sum(axes_to_reduce_drop, /*keepdims=*/false);
}
if (res.shape() != target_shape) {
res = res.Reshape(target_shape);
}
return res;
}
} // namespace
namespace {
Array OutputSlotToArray(const ExecutionResult& r, size_t idx) {
ExecutionResult out;
out.device_buffer = r.device_buffers.at(idx);
out.shape = r.shapes.at(idx);
out.dtype = r.dtypes.at(idx);
out.api = r.api;
return out.AsArray();
}
uint64_t HashTape(const Txpr& tape) {
uint64_t h = 1469598103934665603ULL;
auto mix = [&](uint64_t v) {
h ^= v;
h *= 1099511628211ULL;
};
for (const auto& e : tape.Equations()) {
mix(static_cast<uint64_t>(e.op));
for (const auto& v : e.in_vars) {
mix(v.id);
}
for (const auto& v : e.out_vars) {
mix(v.id);
}
}
return h;
}
struct GradExec {
std::shared_ptr<CachedExecutable> cached;
std::vector<size_t> grad_output_index;
};
GradExec CompileGradProgram(Txpr& tape, const Array& loss_sym, const std::vector<Array>& target_syms, bool include_y) {
VarId cot = tape.AddVar(loss_sym.shape(), loss_sym.Dtype());
tape.MarkInput(cot);
Array cot_sym = Array::Symbolic(cot, loss_sym.shape(), loss_sym.Dtype());
std::unordered_map<uint64_t, Array> gmap;
{
TracerScope scope(tape);
gmap = VjpSweep(loss_sym, target_syms, cot_sym);
}
GradExec ge;
ge.grad_output_index.assign(target_syms.size(), static_cast<size_t>(-1));
size_t next = tape.OutBinders().size();
if (include_y) {
tape.MarkOutput(loss_sym.SymbolicVarId());
next = tape.OutBinders().size();
}
for (size_t i = 0; i < target_syms.size(); ++i) {
auto it = gmap.find(target_syms[i].Id());
if (it != gmap.end() && it->second.IsSymbolic()) {
tape.MarkOutput(it->second.SymbolicVarId());
ge.grad_output_index[i] = next++;
}
}
ge.cached = CompileTxprModule(HashTape(tape), tape);
return ge;
}
std::vector<Array> execute_grads(const GradExec& ge, const std::vector<Array>& forward_inputs, const Array& cotangent,
const std::vector<Array>& target_syms) {
std::vector<Array> inputs = forward_inputs;
inputs.push_back((cotangent.HasData() || cotangent.HasDeviceData()) ? cotangent : cotangent.Eval());
ExecutionResult r = ExecuteJittedExecutable(ge.cached, inputs);
std::vector<Array> grads;
grads.reserve(target_syms.size());
for (size_t i = 0; i < target_syms.size(); ++i) {
if (ge.grad_output_index[i] != static_cast<size_t>(-1)) {
grads.push_back(OutputSlotToArray(r, ge.grad_output_index[i]));
} else {
grads.push_back(Array::Zeros(target_syms[i].shape(), target_syms[i].Dtype()));
}
}
return grads;
}
std::unordered_map<uint64_t, Array> GradViaLowering(const Array& loss, const std::vector<Array>& targets, const Array& cotangent) {
TENZR_CHECK(loss.node() != nullptr, "autograd::grad requires a defined loss graph");
std::unordered_map<uint64_t, VarId> n2v;
Txpr tape = LowerNodeToTxpr(loss.node(), n2v);
Array loss_sym = Array::Symbolic(n2v.at(loss.node()->Id()), loss.shape(), loss.Dtype());
std::vector<Array> target_syms;
target_syms.reserve(targets.size());
for (const auto& t : targets) {
auto it = n2v.find(t.node()->Id());
TENZR_CHECK(it != n2v.end(), "autograd::grad target is not part of the loss graph");
target_syms.push_back(Array::Symbolic(it->second, t.shape(), t.Dtype()));
}
std::vector<Array> params;
for (const auto& p : CollectParameters(loss.node())) {
params.emplace_back(p);
}
GradExec ge = CompileGradProgram(tape, loss_sym, target_syms, /*include_y=*/false);
std::vector<Array> grads = execute_grads(ge, params, cotangent, target_syms);
std::unordered_map<uint64_t, Array> result;
for (size_t i = 0; i < targets.size(); ++i) {
result[targets[i].Id()] = grads[i];
}
return result;
}
} // namespace
std::unordered_map<uint64_t, Array> VjpSweep(const Array& loss, const std::vector<Array>& targets, const Array& cotangent) {
TENZR_TRACE_SCOPE_CAT("autograd", "txpr::grad");
TENZR_CHECK(loss.IsDefined(), "Cannot differentiate undefined loss array");
if (loss.IsSymbolic()) {
auto* tape = txpr::TracerContext::CurrentTape();
TENZR_CHECK(tape != nullptr, "Active tracer tape required for symbolic autograd::grad");
std::vector<txpr::Instruction> fwd_insts = tape->Instructions();
std::unordered_map<uint32_t, Array> adjoints;
adjoints[loss.SymbolicVarId().id] = cotangent;
auto accumulate_grad = [&](txpr::VarId target_var, const Array& grad_val) {
auto it = adjoints.find(target_var.id);
if (it == adjoints.end()) {
adjoints[target_var.id] = grad_val;
} else {
adjoints[target_var.id] = it->second + grad_val;
}
};
auto make_sym_array = [&](txpr::VarId vid) -> Array {
const auto& v = tape->Variable(vid);
return Array::Symbolic(vid, v.shape, v.dtype);
};
for (int64_t i = static_cast<int64_t>(fwd_insts.size()) - 1; i >= 0; --i) {
const auto& eqn = fwd_insts[static_cast<size_t>(i)];
if (eqn.out_vars.empty()) {
continue;
}
auto out_var = eqn.out_vars[0];
auto adj_it = adjoints.find(out_var.id);
if (adj_it == adjoints.end()) {
continue;
}
Array grad_out = adj_it->second;
switch (eqn.op) {
case txpr::OpCode::kAdd: {
accumulate_grad(eqn.in_vars[0], unbroadcast(grad_out, tape->Variable(eqn.in_vars[0]).shape));
accumulate_grad(eqn.in_vars[1], unbroadcast(grad_out, tape->Variable(eqn.in_vars[1]).shape));
break;
}
case txpr::OpCode::kSub: {
accumulate_grad(eqn.in_vars[0], unbroadcast(grad_out, tape->Variable(eqn.in_vars[0]).shape));
accumulate_grad(eqn.in_vars[1], unbroadcast(grad_out.Neg(), tape->Variable(eqn.in_vars[1]).shape));
break;
}
case txpr::OpCode::kMul: {
Array in0 = make_sym_array(eqn.in_vars[0]);
Array in1 = make_sym_array(eqn.in_vars[1]);
accumulate_grad(eqn.in_vars[0], unbroadcast(grad_out * in1, in0.shape()));
accumulate_grad(eqn.in_vars[1], unbroadcast(grad_out * in0, in1.shape()));
break;
}
case txpr::OpCode::kDiv: {
Array in0 = make_sym_array(eqn.in_vars[0]);
Array in1 = make_sym_array(eqn.in_vars[1]);
accumulate_grad(eqn.in_vars[0], unbroadcast(grad_out / in1, in0.shape()));
accumulate_grad(eqn.in_vars[1], unbroadcast((grad_out.Neg() * in0) / (in1 * in1), in1.shape()));
break;
}
case txpr::OpCode::kMatmul: {
Array a = make_sym_array(eqn.in_vars[0]);
Array b = make_sym_array(eqn.in_vars[1]);
accumulate_grad(eqn.in_vars[0], grad_out.Matmul(b.Transpose().Contiguous()));
accumulate_grad(eqn.in_vars[1], a.Transpose().Contiguous().Matmul(grad_out));
break;
}
case txpr::OpCode::kConvolution: {
// NCHW/OIHW conv2d VJP (JAX recipe): grad_X is a strided conv with the
// reversed kernel; grad_W is a conv with the batch/feature dims swapped.
Array input = make_sym_array(eqn.in_vars[0]);
Array kernel = make_sym_array(eqn.in_vars[1]);
const std::vector<int64_t>& strides = eqn.attrs.window_strides;
const std::vector<int64_t>& p_lo = eqn.attrs.edge_padding_low;
const int64_t out_h = grad_out.shape().Dim(2);
const int64_t out_w = grad_out.shape().Dim(3);
std::vector<int64_t> lhs_lo(2), lhs_hi(2), rhs_lo(2), rhs_hi(2);
for (size_t d = 0; d < 2; ++d) {
const int64_t in_d = input.shape().Dim(2 + d);
const int64_t k_d = kernel.shape().Dim(2 + d);
const int64_t s = strides[d];
const int64_t plo = p_lo[d];
const int64_t out_d = (d == 0) ? out_h : out_w;
const int64_t out_dilated = (out_d - 1) * s + 1;
lhs_lo[d] = k_d - plo - 1;
lhs_hi[d] = in_d - out_dilated + plo;
rhs_lo[d] = plo;
rhs_hi[d] = (out_dilated - in_d) + (k_d - plo - 1);
}
Array rev_kernel = plax::Reverse(kernel, {2, 3});
// Contract grad_Y's feature (O) with the reversed kernel's input
// feature by swapping the kernel's O/I dims (equivalent to the
// transposed rhs spec, but expressible with identity specs).
Array rev_kernel_t = rev_kernel.Permute({1, 0, 2, 3});
Array grad_input = plax::ConvGeneral(grad_out, rev_kernel_t, {1, 1}, strides, {1, 1}, lhs_lo, lhs_hi);
// grad_W: contract the batch (N) by transposing X and grad_Y so N
// becomes the feature dim.
Array x_t = input.Permute({1, 0, 2, 3});
Array g_t = grad_out.Permute({1, 0, 2, 3});
Array grad_kernel = plax::ConvGeneral(x_t, g_t, {1, 1}, {1, 1}, strides, rhs_lo, rhs_hi).Permute({1, 0, 2, 3});
accumulate_grad(eqn.in_vars[0], grad_input);
accumulate_grad(eqn.in_vars[1], grad_kernel);
break;
}
case txpr::OpCode::kRelu: {
Array x = make_sym_array(eqn.in_vars[0]);
auto zero = Array::Zeros(x.shape(), x.Dtype());
auto mask = x.GreaterThan(zero).Astype(x.Dtype());
accumulate_grad(eqn.in_vars[0], grad_out * mask);
break;
}
case txpr::OpCode::kSigmoid: {
Array x = make_sym_array(eqn.in_vars[0]);
auto s = x.sigmoid();
auto one = Array::Ones(x.shape(), x.Dtype());
accumulate_grad(eqn.in_vars[0], grad_out * s * (one - s));
break;
}
case txpr::OpCode::kReshape: {
Array in = make_sym_array(eqn.in_vars[0]);
accumulate_grad(eqn.in_vars[0], grad_out.Reshape(in.shape()));
break;
}
case txpr::OpCode::kTranspose: {
const auto& perm = eqn.attrs.permutation;
std::vector<int64_t> inv_perm(perm.size());
for (size_t p = 0; p < perm.size(); ++p) {
inv_perm[static_cast<size_t>(perm[p])] = static_cast<int64_t>(p);
}
accumulate_grad(eqn.in_vars[0], grad_out.Permute(inv_perm));
break;
}
case txpr::OpCode::kBroadcastTo: {
Array in = make_sym_array(eqn.in_vars[0]);
accumulate_grad(eqn.in_vars[0], unbroadcast(grad_out, in.shape()));
break;
}
case txpr::OpCode::kPad: {
Array in = make_sym_array(eqn.in_vars[0]);
const Shape& in_shape = in.shape();
std::vector<Slice> slices(in_shape.Rank());
for (size_t d = 0; d < in_shape.Rank(); ++d) {
const int64_t low = d < eqn.attrs.edge_padding_low.size() ? eqn.attrs.edge_padding_low[d] : 0;
const int64_t interior = d < eqn.attrs.interior_padding.size() ? eqn.attrs.interior_padding[d] : 0;
const int64_t dim = in_shape.Dim(d);
if (interior == 0) {
slices[d] = Slice(low, low + dim, 1);
} else {
slices[d] = Slice(low, low + (dim - 1) * (interior + 1) + 1, interior + 1);
}
}
accumulate_grad(eqn.in_vars[0], grad_out.slice(slices));
break;
}
case txpr::OpCode::kGather: {
Array in = make_sym_array(eqn.in_vars[0]);
Array idx = make_sym_array(eqn.in_vars[1]);
const int64_t axis = eqn.attrs.gather_axis;
Array zero = Array::Zeros(in.shape(), in.Dtype());
if (axis == 0) {
Array idx1 = idx.Reshape(Shape({idx.shape().Dim(0)}));
accumulate_grad(eqn.in_vars[0], plax::ScatterAdd(zero, idx1, grad_out, 0));
} else {
// The symbolic gather keeps the natural [k, ...] order; move k back to
// its operand axis before scattering.
const int64_t rank = static_cast<int64_t>(in.shape().Rank());
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);
}
Array idx1 = idx.Reshape(Shape({idx.shape().Dim(0)}));
accumulate_grad(eqn.in_vars[0], plax::ScatterAdd(zero, idx1, grad_out.Permute(perm), axis));
}
break;
}
case txpr::OpCode::kDynamicSlice: {
Array operand = make_sym_array(eqn.in_vars[0]);
std::vector<Array> starts;
for (size_t vi = 1; vi < eqn.in_vars.size(); ++vi) {
starts.push_back(make_sym_array(eqn.in_vars[vi]));
}
accumulate_grad(eqn.in_vars[0], plax::DynamicUpdateSlice(Array::Zeros(operand.shape(), operand.Dtype()), grad_out, starts));
break;
}
case txpr::OpCode::kDynamicUpdateSlice: {
Array operand = make_sym_array(eqn.in_vars[0]);
Array update = make_sym_array(eqn.in_vars[1]);
std::vector<Array> starts;
for (size_t vi = 2; vi < eqn.in_vars.size(); ++vi) {
starts.push_back(make_sym_array(eqn.in_vars[vi]));
}
accumulate_grad(eqn.in_vars[0], plax::DynamicUpdateSlice(grad_out, Array::Zeros(update.shape(), update.Dtype()), starts));
accumulate_grad(eqn.in_vars[1], plax::DynamicSlice(grad_out, starts, update.shape()));
break;
}
case txpr::OpCode::kScatter: {
TENZR_CHECK(eqn.attrs.scatter_axis == 0, "gradient of plax::scatter with axis != 0 is not implemented");
Array operand = make_sym_array(eqn.in_vars[0]);
Array idx2 = make_sym_array(eqn.in_vars[1]);
Array updates = make_sym_array(eqn.in_vars[2]);
Array idx1 = idx2.Reshape(Shape({idx2.shape().Dim(0)}));
const std::string& mode = eqn.attrs.scatter_update;
Array gathered = plax::Gather(operand, idx1, 0);
Array g_at = plax::Gather(grad_out, idx1, 0);
Array zeros_upd = Array::Zeros(updates.shape(), updates.Dtype());
if (mode == "set") {
accumulate_grad(eqn.in_vars[0], plax::Scatter(grad_out, idx1, zeros_upd, 0, "set"));
} else if (mode == "multiply") {
accumulate_grad(eqn.in_vars[0], plax::Scatter(grad_out, idx1, g_at * updates, 0, "set"));
} else if (mode == "min" || mode == "max") {
Array mask = (mode == "min" ? gathered.LessEqual(updates) : gathered.GreaterEqual(updates)).Astype(operand.Dtype());
accumulate_grad(eqn.in_vars[0], plax::Scatter(grad_out, idx1, g_at * mask, 0, "set"));
} else {
accumulate_grad(eqn.in_vars[0], grad_out);
}
if (mode == "multiply") {
accumulate_grad(eqn.in_vars[2], g_at * gathered);
} else if (mode == "min") {
accumulate_grad(eqn.in_vars[2], g_at * updates.LessThan(gathered).Astype(updates.Dtype()));
} else if (mode == "max") {
accumulate_grad(eqn.in_vars[2], g_at * updates.GreaterThan(gathered).Astype(updates.Dtype()));
} else {
accumulate_grad(eqn.in_vars[2], g_at);
}
break;
}
case txpr::OpCode::kConcatenate: {
const Shape out_shape = tape->Variable(eqn.out_vars[0]).shape;
int64_t axis = eqn.attrs.concat_axis;
if (axis < 0) {
axis += static_cast<int64_t>(out_shape.Rank());
}
int64_t offset = 0;
for (auto vid : eqn.in_vars) {
const int64_t dim = tape->Variable(vid).shape.Dim(static_cast<size_t>(axis));
std::vector<Slice> slices(out_shape.Rank(), Slice::All());
slices[static_cast<size_t>(axis)] = Slice(offset, offset + dim, 1);
accumulate_grad(vid, grad_out.slice(slices));
offset += dim;
}
break;
}
case txpr::OpCode::kArgSort: {
if (!eqn.attrs.sort_return_values) {
// The permutation is piecewise constant in its values: zero gradient.
break;
}
Array in = make_sym_array(eqn.in_vars[0]);
const int64_t axis = eqn.attrs.sort_dimension;
Array perm = plax::Argsort(in, axis, eqn.attrs.sort_descending);
Array inv = plax::Argsort(perm, axis, false);
accumulate_grad(eqn.in_vars[0], plax::TakeAlongAxis(grad_out, inv, axis));
break;
}
case txpr::OpCode::kConvert: {
Array in = make_sym_array(eqn.in_vars[0]);
accumulate_grad(eqn.in_vars[0], grad_out.Astype(in.Dtype()));
break;
}
case txpr::OpCode::kReduceSum: {
Array in = make_sym_array(eqn.in_vars[0]);
if (eqn.attrs.keepdims) {
accumulate_grad(eqn.in_vars[0], grad_out.BroadcastTo(in.shape()));
} else {
const auto& in_shape = in.shape();
int64_t rank = static_cast<int64_t>(in_shape.Rank());
std::vector<int64_t> axes = eqn.attrs.axes;
if (axes.empty()) {
axes.resize(static_cast<size_t>(rank));
std::iota(axes.begin(), axes.end(), 0);
} else {
for (auto& ax : axes) {
if (ax < 0) {
ax += rank;
}
}
}
std::vector<int64_t> exp_dims(static_cast<size_t>(rank));
for (int64_t d = 0; d < rank; ++d) {
bool is_red = std::find(axes.begin(), axes.end(), d) != axes.end();
exp_dims[static_cast<size_t>(d)] = is_red ? 1 : in_shape.Dim(static_cast<size_t>(d));
}
accumulate_grad(eqn.in_vars[0], grad_out.Reshape(Shape(exp_dims)).BroadcastTo(in.shape()));
}
break;
}
case txpr::OpCode::kReduceProd:
case txpr::OpCode::kReduceMax:
case txpr::OpCode::kReduceMin: {
Array in = make_sym_array(eqn.in_vars[0]);
Array red = make_sym_array(eqn.out_vars[0]);
const auto& in_shape = in.shape();
int64_t rank = static_cast<int64_t>(in_shape.Rank());
std::vector<int64_t> axes = eqn.attrs.axes;
if (axes.empty()) {
axes.resize(static_cast<size_t>(rank));
std::iota(axes.begin(), axes.end(), 0);
} else {
for (auto& ax : axes) {
if (ax < 0) {
ax += rank;
}
}
}
std::vector<int64_t> exp_dims(static_cast<size_t>(rank));
for (int64_t d = 0; d < rank; ++d) {
bool is_red = std::find(axes.begin(), axes.end(), d) != axes.end();
exp_dims[static_cast<size_t>(d)] = is_red ? 1 : in_shape.Dim(static_cast<size_t>(d));
}
Array g = eqn.attrs.keepdims ? grad_out : grad_out.Reshape(Shape(exp_dims));
Array red_full = eqn.attrs.keepdims ? red : red.Reshape(Shape(exp_dims));
Array red_b = red_full.BroadcastTo(in_shape);
if (eqn.op == txpr::OpCode::kReduceProd) {
accumulate_grad(eqn.in_vars[0], g.BroadcastTo(in_shape) * (red_b / in));
} else {
Array mask = (eqn.op == txpr::OpCode::kReduceMax ? in.GreaterEqual(red_b) : in.LessEqual(red_b)).Astype(in.Dtype());
accumulate_grad(eqn.in_vars[0], g.BroadcastTo(in_shape) * mask);
}
break;
}
case txpr::OpCode::kSlice: {
Array in = make_sym_array(eqn.in_vars[0]);
const auto& in_shape = in.shape();
size_t rank = in_shape.Rank();
std::vector<int64_t> low(rank, 0);
std::vector<int64_t> high(rank, 0);
std::vector<int64_t> interior(rank, 0);
for (size_t d = 0; d < rank; ++d) {
const int64_t dim = in_shape.Dim(d);
if (d < eqn.attrs.slices.size() && eqn.attrs.slices[d].start.has_value()) {
int64_t st = *eqn.attrs.slices[d].start;
if (st < 0) {
st += dim;
}
int64_t sp = eqn.attrs.slices[d].stop.has_value() ? *eqn.attrs.slices[d].stop : dim;
if (sp < 0) {
sp += dim;
}
const int64_t step = eqn.attrs.slices[d].step;
low[d] = st;
high[d] = dim - sp + (step - 1);
interior[d] = step - 1;
}
}
Array pv = Array::Zeros(Shape({}), in.Dtype());
accumulate_grad(eqn.in_vars[0], plax::Pad(grad_out, pv, low, high, interior));
break;
}
case txpr::OpCode::kCustomOp: {
const std::string& op = eqn.attrs.custom_op_name;
Array in0 = make_sym_array(eqn.in_vars[0]);
if (CustomVJPRegistry::Instance().HasRule(op)) {
const auto& rule = CustomVJPRegistry::Instance().GetRule(op);
Array grad_in = rule.bwd(in0, grad_out);
accumulate_grad(eqn.in_vars[0], grad_in);
} else if (HasDerivativeRule(op)) {
std::vector<Array> inputs;
inputs.reserve(eqn.in_vars.size());
for (auto vid : eqn.in_vars) {
inputs.push_back(make_sym_array(vid));
}
std::vector<Array> grads;
CustomOpVjp(op, inputs, grad_out, grads);
for (size_t k = 0; k < eqn.in_vars.size(); ++k) {
accumulate_grad(eqn.in_vars[k], grads[k]);
}
} else if (IsZeroDerivativeCustomOp(op)) {
// Non-differentiable (comparison / bitwise / shift / rounding): zero gradient.
} else {
TENZR_THROW("gradient rule not implemented for custom op: " + op);
}
break;
}
default:
if (!IsZeroGradientOpcode(eqn.op)) {
TENZR_THROW("gradient rule not implemented for op: " + std::string(OpCodeToString(eqn.op)));
}
break;
}
}
std::unordered_map<uint64_t, Array> result;
for (const auto& t : targets) {
TENZR_CHECK(t.IsDefined(), "Target array in autograd::grad is undefined");
uint64_t tid = t.Id();
auto it = adjoints.find(static_cast<uint32_t>(tid));
if (it != adjoints.end()) {
result[tid] = it->second;
} else {
result[tid] = Array::Zeros(t.shape(), t.Dtype());
}
}
return result;
}
return GradViaLowering(loss, targets, cotangent);
}
std::unordered_map<uint64_t, Array> Grad(const Array& loss, const std::vector<Array>& targets) {
TENZR_CHECK(loss.shape().NumElements() == 1, "autograd::grad requires a scalar loss (received shape: " + loss.shape().ToString() + ")");
return VjpSweep(loss, targets, Array::Ones(loss.shape(), loss.Dtype()));
}
VjpResult Vjp(const std::function<Array(const std::vector<Array>&)>& fn, const std::vector<Array>& primals) {
if (TracerContext::IsActive()) {
Array y = fn(primals);
auto pullback = [y, primals](const Array& cotangent) -> std::vector<Array> {
std::unordered_map<uint64_t, Array> sweep = VjpSweep(y, primals, cotangent);
std::vector<Array> grads;
grads.reserve(primals.size());
for (const auto& p : primals) {
auto it = sweep.find(p.Id());
grads.push_back(it != sweep.end() ? it->second : Array::Zeros(p.shape(), p.Dtype()));
}
return grads;
};
return {y, std::move(pullback)};
}
Txpr tape;
std::vector<Array> syms;
syms.reserve(primals.size());
for (const auto& p : primals) {
VarId v = tape.AddVar(p.shape(), p.Dtype());
tape.MarkInput(v);
syms.push_back(Array::Symbolic(v, p.shape(), p.Dtype()));
}
Array y_sym;
{
TracerScope scope(tape);
y_sym = fn(syms);
}
TENZR_CHECK(y_sym.IsSymbolic(), "txpr::vjp requires a traceable function (output must be symbolic)");
GradExec ge = CompileGradProgram(tape, y_sym, syms, /*include_y=*/true);
std::vector<Array> finputs;
finputs.reserve(primals.size());
for (const auto& p : primals) {
finputs.push_back((p.HasData() || p.HasDeviceData()) ? p : p.Eval());
}
std::vector<Array> y_inputs = finputs;
y_inputs.push_back(Array::Zeros(y_sym.shape(), y_sym.Dtype()));
Array y = OutputSlotToArray(ExecuteJittedExecutable(ge.cached, y_inputs), 0);
auto pullback = [ge, finputs, syms](const Array& cotangent) -> std::vector<Array> { return execute_grads(ge, finputs, cotangent, syms); };
return {y, std::move(pullback)};
}
Array Grad(const Array& loss, const Array& target) {
TENZR_CHECK(target.IsDefined(), "Target array to autograd::grad is undefined");
auto grad_map = Grad(loss, std::vector<Array>{target});
return grad_map.at(target.Id());
}
std::function<Array(const Array&)> Grad(const std::function<Array(const Array&)>& fn) {
return [fn](const Array& x) -> Array {
auto traced = [fn](const std::vector<Array>& ps) -> Array { return fn(ps[0]); };
auto [y, pullback] = Vjp(traced, {x});
TENZR_CHECK(y.shape().NumElements() == 1, "autograd::grad requires a scalar loss");
return pullback(Array::Ones(y.shape(), y.Dtype()))[0];
};
}
std::function<ValueAndGradResult(const Array&)> ValueAndGrad(const std::function<Array(const Array&)>& fn) {
return [fn](const Array& x) -> ValueAndGradResult {
auto traced = [fn](const std::vector<Array>& ps) -> Array { return fn(ps[0]); };
auto [y, pullback] = Vjp(traced, {x});
return {y, pullback(Array::Ones(y.shape(), y.Dtype()))[0]};
};
}
} // namespace tenzr::txpr