// 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/jvp.h"
#include <algorithm>
#include <memory>
#include <numeric>
#include <stdexcept>
#include <string>
#include <unordered_map>
#include <vector>
#include "tenzr/common/error.h"
#include "tenzr/plax/plax.h"
#include "tenzr/txpr/custom_op_deriv.h"
#include "tenzr/txpr/tracer.h"
#include "tenzr/txpr/txpr.h"
namespace tenzr::txpr {
namespace {
Array ConstantPrimal(const AttributePayload& attrs, const Shape& shape, DType dtype) {
if (!attrs.tensor_data.empty()) {
if (dtype == DType::kFloat64) {
Array a = Array::Zeros(shape, DType::kFloat64);
a.SetData<double>(attrs.tensor_data);
return a;
}
std::vector<float> f(attrs.tensor_data.begin(), attrs.tensor_data.end());
Array a = Array::Zeros(shape, DType::kFloat32);
a.SetData<float>(f);
return a;
}
if (!attrs.int_tensor_data.empty()) {
const auto& d = attrs.int_tensor_data;
switch (dtype) {
case DType::kBool: {
std::vector<bool> v;
v.reserve(d.size());
for (int64_t x : d) {
v.push_back(x != 0);
}
Array a = Array::Zeros(shape, DType::kBool);
a.SetData<bool>(v);
return a;
}
case DType::kInt8: {
std::vector<int8_t> v(d.begin(), d.end());
Array a = Array::Zeros(shape, DType::kInt8);
a.SetData<int8_t>(v);
return a;
}
case DType::kInt16: {
std::vector<int16_t> v(d.begin(), d.end());
Array a = Array::Zeros(shape, DType::kInt16);
a.SetData<int16_t>(v);
return a;
}
case DType::kInt32: {
std::vector<int32_t> v(d.begin(), d.end());
Array a = Array::Zeros(shape, DType::kInt32);
a.SetData<int32_t>(v);
return a;
}
case DType::kUInt8: {
std::vector<uint8_t> v(d.begin(), d.end());
Array a = Array::Zeros(shape, DType::kUInt8);
a.SetData<uint8_t>(v);
return a;
}
case DType::kUInt16: {
std::vector<uint16_t> v(d.begin(), d.end());
Array a = Array::Zeros(shape, DType::kUInt16);
a.SetData<uint16_t>(v);
return a;
}
case DType::kUInt32: {
std::vector<uint32_t> v(d.begin(), d.end());
Array a = Array::Zeros(shape, DType::kUInt32);
a.SetData<uint32_t>(v);
return a;
}
case DType::kUInt64: {
std::vector<uint64_t> v(d.begin(), d.end());
Array a = Array::Zeros(shape, DType::kUInt64);
a.SetData<uint64_t>(v);
return a;
}
default: {
Array a = Array::Zeros(shape, DType::kInt64);
a.SetData<int64_t>(d);
return a;
}
}
}
return Array::Full(shape, attrs.scalar_value, dtype);
}
} // namespace
JvpResult Jvp(const std::function<Array(const Array&)>& fn, const Array& primal, const Array& tangent) {
TENZR_CHECK(primal.shape() == tangent.shape(), "jvp primal and tangent shapes must match");
TENZR_CHECK(primal.Dtype() == tangent.Dtype(), "jvp primal and tangent dtypes must match");
Txpr txpr;
VarId in_vid = txpr.AddVar(primal.shape(), primal.Dtype(), "primal_in");
txpr.MarkInput(in_vid);
Array sym_in = Array::Symbolic(in_vid, primal.shape(), primal.Dtype());
Array sym_out;
{
TracerScope scope(txpr);
sym_out = fn(sym_in);
}
txpr.MarkOutput(sym_out.SymbolicVarId());
// Forward primal and tangent propagation along Txpr equations
std::unordered_map<uint32_t, Array> primals;
std::unordered_map<uint32_t, Array> tangents;
primals[in_vid.id] = primal;
tangents[in_vid.id] = tangent;
auto get_primal = [&](VarId vid) -> Array {
auto it = primals.find(vid.id);
TENZR_CHECK(it != primals.end(), "jvp missing primal value for variable " + std::to_string(vid.id));
return it->second;
};
auto get_tangent = [&](VarId vid) -> Array {
auto it = tangents.find(vid.id);
if (it != tangents.end()) {
return it->second;
}
const auto& vdef = txpr.Variable(vid);
return Array::Zeros(vdef.shape, vdef.dtype);
};
for (const auto& eqn : txpr.Equations()) {
VarId out_vid = eqn.out_vars[0];
const auto& out_vdef = txpr.Variable(out_vid);
switch (eqn.op) {
case OpCode::kAdd: {
Array p0 = get_primal(eqn.in_vars[0]);
Array p1 = get_primal(eqn.in_vars[1]);
Array t0 = get_tangent(eqn.in_vars[0]);
Array t1 = get_tangent(eqn.in_vars[1]);
primals[out_vid.id] = p0 + p1;
tangents[out_vid.id] = t0 + t1;
break;
}
case OpCode::kSub: {
Array p0 = get_primal(eqn.in_vars[0]);
Array p1 = get_primal(eqn.in_vars[1]);
Array t0 = get_tangent(eqn.in_vars[0]);
Array t1 = get_tangent(eqn.in_vars[1]);
primals[out_vid.id] = p0 - p1;
tangents[out_vid.id] = t0 - t1;
break;
}
case OpCode::kMul: {
Array p0 = get_primal(eqn.in_vars[0]);
Array p1 = get_primal(eqn.in_vars[1]);
Array t0 = get_tangent(eqn.in_vars[0]);
Array t1 = get_tangent(eqn.in_vars[1]);
primals[out_vid.id] = p0 * p1;
tangents[out_vid.id] = t0 * p1 + p0 * t1;
break;
}
case OpCode::kDiv: {
Array p0 = get_primal(eqn.in_vars[0]);
Array p1 = get_primal(eqn.in_vars[1]);
Array t0 = get_tangent(eqn.in_vars[0]);
Array t1 = get_tangent(eqn.in_vars[1]);
primals[out_vid.id] = p0 / p1;
tangents[out_vid.id] = (t0 * p1 - p0 * t1) / (p1 * p1);
break;
}
case OpCode::kMatmul: {
Array p0 = get_primal(eqn.in_vars[0]);
Array p1 = get_primal(eqn.in_vars[1]);
Array t0 = get_tangent(eqn.in_vars[0]);
Array t1 = get_tangent(eqn.in_vars[1]);
primals[out_vid.id] = p0.Matmul(p1);
tangents[out_vid.id] = t0.Matmul(p1) + p0.Matmul(t1);
break;
}
case OpCode::kConvolution: {
Array p0 = get_primal(eqn.in_vars[0]);
Array p1 = get_primal(eqn.in_vars[1]);
Array t0 = get_tangent(eqn.in_vars[0]);
Array t1 = get_tangent(eqn.in_vars[1]);
const auto& st = eqn.attrs.window_strides;
const auto& ld = eqn.attrs.lhs_dilation;
const auto& rd = eqn.attrs.rhs_dilation;
const auto& plo = eqn.attrs.edge_padding_low;
const auto& phi = eqn.attrs.edge_padding_high;
const auto& ls = eqn.attrs.conv_lhs_spec;
const auto& rs = eqn.attrs.conv_rhs_spec;
const auto& os = eqn.attrs.conv_out_spec;
primals[out_vid.id] = plax::ConvGeneral(p0, p1, st, ld, rd, plo, phi, ls, rs, os);
tangents[out_vid.id] = plax::ConvGeneral(t0, p1, st, ld, rd, plo, phi, ls, rs, os) + plax::ConvGeneral(p0, t1, st, ld, rd, plo, phi, ls, rs, os);
break;
}
case OpCode::kReshape: {
Array p0 = get_primal(eqn.in_vars[0]);
Array t0 = get_tangent(eqn.in_vars[0]);
primals[out_vid.id] = p0.Reshape(out_vdef.shape);
tangents[out_vid.id] = t0.Reshape(out_vdef.shape);
break;
}
case OpCode::kTranspose: {
Array p0 = get_primal(eqn.in_vars[0]);
Array t0 = get_tangent(eqn.in_vars[0]);
primals[out_vid.id] = p0.Permute(eqn.attrs.permutation);
tangents[out_vid.id] = t0.Permute(eqn.attrs.permutation);
break;
}
case OpCode::kBroadcastTo: {
Array p0 = get_primal(eqn.in_vars[0]);
Array t0 = get_tangent(eqn.in_vars[0]);
primals[out_vid.id] = p0.BroadcastTo(out_vdef.shape);
tangents[out_vid.id] = t0.BroadcastTo(out_vdef.shape);
break;
}
case OpCode::kSlice: {
Array p0 = get_primal(eqn.in_vars[0]);
Array t0 = get_tangent(eqn.in_vars[0]);
primals[out_vid.id] = p0.slice(eqn.attrs.slices);
tangents[out_vid.id] = t0.slice(eqn.attrs.slices);
break;
}
case OpCode::kReduceSum: {
Array p0 = get_primal(eqn.in_vars[0]);
Array t0 = get_tangent(eqn.in_vars[0]);
auto axes_opt = eqn.attrs.axes.empty() ? std::nullopt : std::optional(eqn.attrs.axes);
primals[out_vid.id] = p0.Sum(axes_opt, eqn.attrs.keepdims);
tangents[out_vid.id] = t0.Sum(axes_opt, eqn.attrs.keepdims);
break;
}
case OpCode::kReduceProd:
case OpCode::kReduceMax:
case OpCode::kReduceMin: {
Array p0 = get_primal(eqn.in_vars[0]);
Array t0 = get_tangent(eqn.in_vars[0]);
auto axes_opt = eqn.attrs.axes.empty() ? std::nullopt : std::optional(eqn.attrs.axes);
const int64_t rank = static_cast<int64_t>(p0.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);
}
std::vector<int64_t> red_dims(static_cast<size_t>(rank));
for (int64_t d = 0; d < rank; ++d) {
red_dims[static_cast<size_t>(d)] = (std::find(axes.begin(), axes.end(), d) != axes.end()) ? 1 : p0.shape().Dim(static_cast<size_t>(d));
}
Array prim = (eqn.op == OpCode::kReduceProd) ? p0.Prod(axes_opt, eqn.attrs.keepdims)
: (eqn.op == OpCode::kReduceMax) ? p0.Max(axes_opt, eqn.attrs.keepdims)
: p0.Min(axes_opt, eqn.attrs.keepdims);
primals[out_vid.id] = prim;
Array prim_b = (eqn.attrs.keepdims ? prim : prim.Reshape(Shape(red_dims))).BroadcastTo(p0.shape());
if (eqn.op == OpCode::kReduceProd) {
tangents[out_vid.id] = (t0 * (prim_b / p0)).Sum(axes_opt, eqn.attrs.keepdims);
} else {
Array mask = (eqn.op == OpCode::kReduceMax ? p0.GreaterEqual(prim_b) : p0.LessEqual(prim_b)).Astype(p0.Dtype());
tangents[out_vid.id] = (t0 * mask).Sum(axes_opt, eqn.attrs.keepdims);
}
break;
}
case OpCode::kRelu: {
Array p0 = get_primal(eqn.in_vars[0]);
Array t0 = get_tangent(eqn.in_vars[0]);
primals[out_vid.id] = p0.Relu();
Array zero = Array::Zeros(p0.shape(), p0.Dtype());
Array mask = p0.GreaterThan(zero).Astype(p0.Dtype());
tangents[out_vid.id] = t0 * mask;
break;
}
case OpCode::kSigmoid: {
Array p0 = get_primal(eqn.in_vars[0]);
Array t0 = get_tangent(eqn.in_vars[0]);
Array s = p0.sigmoid();
primals[out_vid.id] = s;
Array one = Array::Ones(p0.shape(), p0.Dtype());
tangents[out_vid.id] = t0 * s * (one - s);
break;
}
case OpCode::kConstant: {
primals[out_vid.id] = ConstantPrimal(eqn.attrs, out_vdef.shape, out_vdef.dtype);
tangents[out_vid.id] = Array::Zeros(out_vdef.shape, out_vdef.dtype);
break;
}
case OpCode::kCustomOp: {
const std::string& cop = eqn.attrs.custom_op_name;
std::vector<Array> cps;
std::vector<Array> cts;
cps.reserve(eqn.in_vars.size());
cts.reserve(eqn.in_vars.size());
for (const auto& vid : eqn.in_vars) {
cps.push_back(get_primal(vid));
cts.push_back(get_tangent(vid));
}
if (HasJvpRule(cop)) {
primals[out_vid.id] = Array::CustomOp(cop, cps);
tangents[out_vid.id] = CustomOpJvp(cop, cps, cts);
} else if (IsZeroDerivativeCustomOp(cop)) {
primals[out_vid.id] = Array::CustomOp(cop, std::move(cps));
tangents[out_vid.id] = Array::Zeros(out_vdef.shape, out_vdef.dtype);
} else {
TENZR_THROW("jvp rule not implemented for custom op: " + cop);
}
break;
}
case OpCode::kConvert: {
Array p0 = get_primal(eqn.in_vars[0]);
Array t0 = get_tangent(eqn.in_vars[0]);
primals[out_vid.id] = p0.Astype(out_vdef.dtype);
tangents[out_vid.id] = t0.Astype(out_vdef.dtype);
break;
}
case OpCode::kPad: {
Array p0 = get_primal(eqn.in_vars[0]);
Array t0 = get_tangent(eqn.in_vars[0]);
Array pv = get_primal(eqn.in_vars[1]);
Array tv = get_tangent(eqn.in_vars[1]);
primals[out_vid.id] = plax::Pad(p0, pv, eqn.attrs.edge_padding_low, eqn.attrs.edge_padding_high, eqn.attrs.interior_padding);
tangents[out_vid.id] = plax::Pad(t0, tv, eqn.attrs.edge_padding_low, eqn.attrs.edge_padding_high, eqn.attrs.interior_padding);
break;
}
case OpCode::kGather: {
Array p0 = get_primal(eqn.in_vars[0]);
Array t0 = get_tangent(eqn.in_vars[0]);
Array pidx = get_primal(eqn.in_vars[1]);
Array idx1 = pidx.Reshape(Shape({pidx.shape().Dim(0)}));
primals[out_vid.id] = plax::Gather(p0, idx1, eqn.attrs.gather_axis);
tangents[out_vid.id] = plax::Gather(t0, idx1, eqn.attrs.gather_axis);
break;
}
case OpCode::kScatter: {
Array p0 = get_primal(eqn.in_vars[0]);
Array t0 = get_tangent(eqn.in_vars[0]);
Array pidx2 = get_primal(eqn.in_vars[1]);
Array pu = get_primal(eqn.in_vars[2]);
Array tu = get_tangent(eqn.in_vars[2]);
Array idx1 = pidx2.Reshape(Shape({pidx2.shape().Dim(0)}));
const std::string& mode = eqn.attrs.scatter_update;
primals[out_vid.id] = plax::Scatter(p0, idx1, pu, eqn.attrs.scatter_axis, mode);
if (mode == "set" || mode == "add") {
tangents[out_vid.id] = plax::Scatter(t0, idx1, tu, eqn.attrs.scatter_axis, mode);
} else if (mode == "multiply") {
Array gp = plax::Gather(p0, idx1, eqn.attrs.scatter_axis);
Array gt = plax::Gather(t0, idx1, eqn.attrs.scatter_axis);
tangents[out_vid.id] = plax::Scatter(t0, idx1, gt * pu + gp * tu, eqn.attrs.scatter_axis, "set");
} else {
Array gp = plax::Gather(p0, idx1, eqn.attrs.scatter_axis);
Array gt = plax::Gather(t0, idx1, eqn.attrs.scatter_axis);
Array m0 = (mode == "min" ? gp.LessEqual(pu) : gp.GreaterEqual(pu)).Astype(p0.Dtype());
Array m1 = (mode == "min" ? pu.LessThan(gp) : pu.GreaterThan(gp)).Astype(p0.Dtype());
tangents[out_vid.id] = plax::Scatter(t0, idx1, gt * m0 + tu * m1, eqn.attrs.scatter_axis, "set");
}
break;
}
case OpCode::kConcatenate: {
std::vector<Array> pps;
std::vector<Array> tts;
pps.reserve(eqn.in_vars.size());
tts.reserve(eqn.in_vars.size());
for (auto vid : eqn.in_vars) {
pps.push_back(get_primal(vid));
tts.push_back(get_tangent(vid));
}
primals[out_vid.id] = plax::Concatenate(pps, eqn.attrs.concat_axis);
tangents[out_vid.id] = plax::Concatenate(tts, eqn.attrs.concat_axis);
break;
}
case OpCode::kArgSort: {
Array p0 = get_primal(eqn.in_vars[0]);
const int64_t axis = eqn.attrs.sort_dimension;
if (eqn.attrs.sort_return_values) {
Array perm = plax::Argsort(p0, axis, eqn.attrs.sort_descending);
primals[out_vid.id] = plax::Sort(p0, axis, eqn.attrs.sort_descending);
tangents[out_vid.id] = plax::TakeAlongAxis(get_tangent(eqn.in_vars[0]), perm, axis);
break;
}
primals[out_vid.id] = plax::Argsort(p0, axis, eqn.attrs.sort_descending);
tangents[out_vid.id] = Array::Zeros(out_vdef.shape, out_vdef.dtype);
break;
}
default: {
TENZR_THROW("jvp rule not implemented for op: " + std::string(OpCodeToString(eqn.op)));
}
}
}
return {get_primal(sym_out.SymbolicVarId()), get_tangent(sym_out.SymbolicVarId())};
}
LinearizeResult Linearize(const std::function<Array(const Array&)>& fn, const Array& primal) {
// 1. Trace and evaluate forward pass to capture all intermediate primals
Txpr txpr;
VarId in_vid = txpr.AddVar(primal.shape(), primal.Dtype(), "primal_in");
txpr.MarkInput(in_vid);
Array sym_in = Array::Symbolic(in_vid, primal.shape(), primal.Dtype());
Array sym_out;
{
TracerScope scope(txpr);
sym_out = fn(sym_in);
}
txpr.MarkOutput(sym_out.SymbolicVarId());
// Forward evaluation of primals
auto captured_primals = std::make_shared<std::unordered_map<uint32_t, Array>>();
(*captured_primals)[in_vid.id] = primal;
auto get_p = [captured_primals](VarId vid) -> Array { return captured_primals->at(vid.id); };
for (const auto& eqn : txpr.Equations()) {
VarId out_vid = eqn.out_vars[0];
const auto& out_vdef = txpr.Variable(out_vid);
switch (eqn.op) {
case OpCode::kAdd:
(*captured_primals)[out_vid.id] = get_p(eqn.in_vars[0]) + get_p(eqn.in_vars[1]);
break;
case OpCode::kSub:
(*captured_primals)[out_vid.id] = get_p(eqn.in_vars[0]) - get_p(eqn.in_vars[1]);
break;
case OpCode::kMul:
(*captured_primals)[out_vid.id] = get_p(eqn.in_vars[0]) * get_p(eqn.in_vars[1]);
break;
case OpCode::kDiv:
(*captured_primals)[out_vid.id] = get_p(eqn.in_vars[0]) / get_p(eqn.in_vars[1]);
break;
case OpCode::kMatmul:
(*captured_primals)[out_vid.id] = get_p(eqn.in_vars[0]).Matmul(get_p(eqn.in_vars[1]));
break;
case OpCode::kReshape:
(*captured_primals)[out_vid.id] = get_p(eqn.in_vars[0]).Reshape(out_vdef.shape);
break;
case OpCode::kTranspose:
(*captured_primals)[out_vid.id] = get_p(eqn.in_vars[0]).Permute(eqn.attrs.permutation);
break;
case OpCode::kBroadcastTo:
(*captured_primals)[out_vid.id] = get_p(eqn.in_vars[0]).BroadcastTo(out_vdef.shape);
break;
case OpCode::kSlice:
(*captured_primals)[out_vid.id] = get_p(eqn.in_vars[0]).slice(eqn.attrs.slices);
break;
case OpCode::kReduceSum: {
auto axes_opt = eqn.attrs.axes.empty() ? std::nullopt : std::optional(eqn.attrs.axes);
(*captured_primals)[out_vid.id] = get_p(eqn.in_vars[0]).Sum(axes_opt, eqn.attrs.keepdims);
break;
}
case OpCode::kReduceProd:
case OpCode::kReduceMax:
case OpCode::kReduceMin: {
auto axes_opt = eqn.attrs.axes.empty() ? std::nullopt : std::optional(eqn.attrs.axes);
Array in = get_p(eqn.in_vars[0]);
if (eqn.op == OpCode::kReduceProd) {
(*captured_primals)[out_vid.id] = in.Prod(axes_opt, eqn.attrs.keepdims);
} else if (eqn.op == OpCode::kReduceMax) {
(*captured_primals)[out_vid.id] = in.Max(axes_opt, eqn.attrs.keepdims);
} else {
(*captured_primals)[out_vid.id] = in.Min(axes_opt, eqn.attrs.keepdims);
}
break;
}
case OpCode::kRelu:
(*captured_primals)[out_vid.id] = get_p(eqn.in_vars[0]).Relu();
break;
case OpCode::kSigmoid:
(*captured_primals)[out_vid.id] = get_p(eqn.in_vars[0]).sigmoid();
break;
case OpCode::kConstant:
(*captured_primals)[out_vid.id] = ConstantPrimal(eqn.attrs, out_vdef.shape, out_vdef.dtype);
break;
case OpCode::kCustomOp: {
const std::string& cop = eqn.attrs.custom_op_name;
std::vector<Array> cps;
cps.reserve(eqn.in_vars.size());
for (const auto& vid : eqn.in_vars) {
cps.push_back(get_p(vid));
}
(*captured_primals)[out_vid.id] = Array::CustomOp(cop, std::move(cps));
break;
}
case OpCode::kConvert:
(*captured_primals)[out_vid.id] = get_p(eqn.in_vars[0]).Astype(out_vdef.dtype);
break;
case OpCode::kPad:
(*captured_primals)[out_vid.id] =
plax::Pad(get_p(eqn.in_vars[0]), get_p(eqn.in_vars[1]), eqn.attrs.edge_padding_low, eqn.attrs.edge_padding_high, eqn.attrs.interior_padding);
break;
case OpCode::kGather: {
Array pidx = get_p(eqn.in_vars[1]);
(*captured_primals)[out_vid.id] = plax::Gather(get_p(eqn.in_vars[0]), pidx.Reshape(Shape({pidx.shape().Dim(0)})), eqn.attrs.gather_axis);
break;
}
case OpCode::kScatter: {
Array pidx = get_p(eqn.in_vars[1]);
(*captured_primals)[out_vid.id] = plax::Scatter(get_p(eqn.in_vars[0]), pidx.Reshape(Shape({pidx.shape().Dim(0)})), get_p(eqn.in_vars[2]),
eqn.attrs.scatter_axis, eqn.attrs.scatter_update);
break;
}
case OpCode::kConcatenate: {
std::vector<Array> pps;
pps.reserve(eqn.in_vars.size());
for (auto vid : eqn.in_vars) {
pps.push_back(get_p(vid));
}
(*captured_primals)[out_vid.id] = plax::Concatenate(pps, eqn.attrs.concat_axis);
break;
}
case OpCode::kArgSort: {
Array p0 = get_p(eqn.in_vars[0]);
(*captured_primals)[out_vid.id] = eqn.attrs.sort_return_values ? plax::Sort(p0, eqn.attrs.sort_dimension, eqn.attrs.sort_descending)
: plax::Argsort(p0, eqn.attrs.sort_dimension, eqn.attrs.sort_descending);
break;
}
default:
TENZR_THROW("jvp primal rule not implemented for op: " + std::string(OpCodeToString(eqn.op)));
}
}
Array primal_out = get_p(sym_out.SymbolicVarId());
// 2. Return primal_out and linearized tangent evaluator
auto shared_txpr = std::make_shared<Txpr>(std::move(txpr));
VarId out_var_id = sym_out.SymbolicVarId();
auto lin_fn = [shared_txpr, captured_primals, in_vid, out_var_id](const Array& tangent) -> Array {
std::unordered_map<uint32_t, Array> tangents;
tangents[in_vid.id] = tangent;
auto get_t = [&](VarId vid) -> Array {
auto it = tangents.find(vid.id);
if (it != tangents.end()) {
return it->second;
}
const auto& vdef = shared_txpr->Variable(vid);
return Array::Zeros(vdef.shape, vdef.dtype);
};
auto get_p_static = [&](VarId vid) -> Array { return captured_primals->at(vid.id); };
for (const auto& eqn : shared_txpr->Equations()) {
VarId out_vid = eqn.out_vars[0];
const auto& out_vdef = shared_txpr->Variable(out_vid);
switch (eqn.op) {
case OpCode::kAdd:
tangents[out_vid.id] = get_t(eqn.in_vars[0]) + get_t(eqn.in_vars[1]);
break;
case OpCode::kSub:
tangents[out_vid.id] = get_t(eqn.in_vars[0]) - get_t(eqn.in_vars[1]);
break;
case OpCode::kMul: {
Array p0 = get_p_static(eqn.in_vars[0]);
Array p1 = get_p_static(eqn.in_vars[1]);
Array t0 = get_t(eqn.in_vars[0]);
Array t1 = get_t(eqn.in_vars[1]);
tangents[out_vid.id] = t0 * p1 + p0 * t1;
break;
}
case OpCode::kDiv: {
Array p0 = get_p_static(eqn.in_vars[0]);
Array p1 = get_p_static(eqn.in_vars[1]);
Array t0 = get_t(eqn.in_vars[0]);
Array t1 = get_t(eqn.in_vars[1]);
tangents[out_vid.id] = (t0 * p1 - p0 * t1) / (p1 * p1);
break;
}
case OpCode::kMatmul: {
Array p0 = get_p_static(eqn.in_vars[0]);
Array p1 = get_p_static(eqn.in_vars[1]);
Array t0 = get_t(eqn.in_vars[0]);
Array t1 = get_t(eqn.in_vars[1]);
tangents[out_vid.id] = t0.Matmul(p1) + p0.Matmul(t1);
break;
}
case OpCode::kReshape:
tangents[out_vid.id] = get_t(eqn.in_vars[0]).Reshape(out_vdef.shape);
break;
case OpCode::kTranspose:
tangents[out_vid.id] = get_t(eqn.in_vars[0]).Permute(eqn.attrs.permutation);
break;
case OpCode::kBroadcastTo:
tangents[out_vid.id] = get_t(eqn.in_vars[0]).BroadcastTo(out_vdef.shape);
break;
case OpCode::kSlice:
tangents[out_vid.id] = get_t(eqn.in_vars[0]).slice(eqn.attrs.slices);
break;
case OpCode::kReduceSum: {
auto axes_opt = eqn.attrs.axes.empty() ? std::nullopt : std::optional(eqn.attrs.axes);
tangents[out_vid.id] = get_t(eqn.in_vars[0]).Sum(axes_opt, eqn.attrs.keepdims);
break;
}
case OpCode::kRelu: {
Array p0 = get_p_static(eqn.in_vars[0]);
Array zero = Array::Zeros(p0.shape(), p0.Dtype());
Array mask = p0.GreaterThan(zero).Astype(p0.Dtype());
tangents[out_vid.id] = get_t(eqn.in_vars[0]) * mask;
break;
}
case OpCode::kSigmoid: {
Array s = get_p_static(out_vid);
Array one = Array::Ones(s.shape(), s.Dtype());
tangents[out_vid.id] = get_t(eqn.in_vars[0]) * s * (one - s);
break;
}
case OpCode::kConstant:
tangents[out_vid.id] = Array::Zeros(out_vdef.shape, out_vdef.dtype);
break;
case OpCode::kCustomOp: {
const std::string& cop = eqn.attrs.custom_op_name;
std::vector<Array> cps;
std::vector<Array> cts;
cps.reserve(eqn.in_vars.size());
cts.reserve(eqn.in_vars.size());
for (const auto& vid : eqn.in_vars) {
cps.push_back(get_p_static(vid));
cts.push_back(get_t(vid));
}
if (HasJvpRule(cop)) {
tangents[out_vid.id] = CustomOpJvp(cop, cps, cts);
} else {
tangents[out_vid.id] = cts[0];
}
break;
}
default:
tangents[out_vid.id] = get_t(eqn.in_vars[0]);
break;
}
}
return get_t(out_var_id);
};
return {primal_out, lin_fn};
}
} // namespace tenzr::txpr