// 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/custom_op_deriv.h"
#include <string>
#include <vector>
#include "tenzr/common/error.h"
#include "tenzr/plax/plax.h"
namespace tenzr::txpr {
namespace {
Array unbroadcast(const Array& g, const Shape& target) {
if (g.shape() == target) {
return g;
}
if (target.NumElements() == 1) {
return g.Sum(std::nullopt, false).Reshape(target);
}
std::vector<int64_t> axes;
int64_t extra = static_cast<int64_t>(g.shape().Rank()) - static_cast<int64_t>(target.Rank());
for (int64_t d = 0; d < extra; ++d) {
axes.push_back(d);
}
int64_t g_rank = static_cast<int64_t>(g.shape().Rank());
int64_t t_rank = static_cast<int64_t>(target.Rank());
for (int64_t d = 0; d < t_rank; ++d) {
if (target.Dim(static_cast<size_t>(d)) == 1 && g.shape().Dim(static_cast<size_t>(g_rank - t_rank + d)) != 1) {
axes.push_back(g_rank - t_rank + d);
}
}
Array r = g;
if (!axes.empty()) {
r = r.Sum(axes, true);
}
return r.Reshape(target);
}
void vjp(std::string_view op, const std::vector<Array>& in, const Array& g, std::vector<Array>& out) {
const Array& a = in[0];
if (op == "tanh") {
Array t = a.tanh();
out.push_back(g * (Array::Ones(a.shape(), a.Dtype()) - t * t));
} else if (op == "neg") {
out.push_back(g.Neg());
} else if (op == "exp") {
out.push_back(g * a.Exp());
} else if (op == "log") {
out.push_back(g / a);
} else if (op == "sqrt") {
out.push_back(g / (Array::Full(a.shape(), 2.0, a.Dtype()) * a.Sqrt()));
} else if (op == "sin") {
out.push_back(g * a.Cos());
} else if (op == "cos") {
out.push_back((g * a.Sin()).Neg());
} else if (op == "abs") {
Array mask = a.GreaterThan(Array::Zeros(a.shape(), a.Dtype())).Astype(a.Dtype());
Array pmask = Array::Ones(a.shape(), a.Dtype()) - mask;
out.push_back(g * (mask - pmask));
} else if (op == "rsqrt") {
Array r = a.Rsqrt();
out.push_back((g * Array::Full(a.shape(), 0.5, a.Dtype()) * r).Neg() / a);
} else if (op == "cbrt") {
Array c = a.Cbrt();
out.push_back(g / (Array::Full(a.shape(), 3.0, a.Dtype()) * c * c));
} else if (op == "expm1") {
out.push_back(g * a.Exp());
} else if (op == "log1p") {
out.push_back(g / (Array::Ones(a.shape(), a.Dtype()) + a));
} else if (op == "tan") {
Array t = a.Tan();
out.push_back(g * (Array::Ones(a.shape(), a.Dtype()) + t * t));
} else if (op == "pow" && in.size() == 2) {
const Array& b = in[1];
out.push_back(unbroadcast(g * b * a.Pow(b - Array::Full(b.shape(), 1.0, b.Dtype())), a.shape()));
out.push_back(unbroadcast(g * a.Pow(b) * a.Log(), b.shape()));
} else if (op == "maximum" && in.size() == 2) {
const Array& b = in[1];
out.push_back(unbroadcast(g * a.GreaterEqual(b).Astype(a.Dtype()), a.shape()));
out.push_back(unbroadcast(g * b.GreaterThan(a).Astype(b.Dtype()), b.shape()));
} else if (op == "minimum" && in.size() == 2) {
const Array& b = in[1];
out.push_back(unbroadcast(g * a.LessEqual(b).Astype(a.Dtype()), a.shape()));
out.push_back(unbroadcast(g * b.LessThan(a).Astype(b.Dtype()), b.shape()));
} else if (op == "atan2" && in.size() == 2) {
const Array& b = in[1];
Array denom = a * a + b * b;
out.push_back(unbroadcast(g * b / denom, a.shape()));
out.push_back(unbroadcast((g * a).Neg() / denom, b.shape()));
} else if (op == "remainder" && in.size() == 2) {
const Array& b = in[1];
out.push_back(unbroadcast(g, a.shape()));
out.push_back(unbroadcast(g * (a / b).Floor().Neg(), b.shape()));
} else if (op == "select" && in.size() == 3) {
const Array& pred = in[0];
Array z = Array::Zeros(g.shape(), g.Dtype());
out.push_back(Array::Zeros(pred.shape(), pred.Dtype()));
out.push_back(unbroadcast(plax::Select(pred, g, z), in[1].shape()));
out.push_back(unbroadcast(plax::Select(pred, z, g), in[2].shape()));
} else if (op == "clamp" && in.size() == 3) {
const Array& lo = in[0];
const Array& x = in[1];
const Array& hi = in[2];
Array in_range = x.GreaterEqual(lo).Astype(x.Dtype()) * x.LessEqual(hi).Astype(x.Dtype());
out.push_back(g * x.LessThan(lo).Astype(x.Dtype()));
out.push_back(g * in_range);
out.push_back(g * x.GreaterThan(hi).Astype(x.Dtype()));
} else {
TENZR_THROW("gradient rule not implemented for custom op: " + std::string(op));
}
}
Array jvp(std::string_view op, const std::vector<Array>& p, const std::vector<Array>& t) {
const Array& a = p[0];
if (op == "tanh") {
Array th = a.tanh();
return t[0] * (Array::Ones(a.shape(), a.Dtype()) - th * th);
}
if (op == "neg") {
return -t[0];
}
if (op == "exp") {
return t[0] * a.Exp();
}
if (op == "log") {
return t[0] / a;
}
if (op == "sqrt") {
return t[0] / (Array::Full(a.shape(), 2.0, a.Dtype()) * a.Sqrt());
}
if (op == "sin") {
return t[0] * a.Cos();
}
if (op == "cos") {
return -(t[0] * a.Sin());
}
if (op == "abs") {
return t[0] * a.Sign();
}
if (op == "rsqrt") {
Array r = a.Rsqrt();
return Array::Full(a.shape(), -0.5, a.Dtype()) * t[0] * r * r * r;
}
if (op == "cbrt") {
Array c = a.Cbrt();
return t[0] / (Array::Full(a.shape(), 3.0, a.Dtype()) * c * c);
}
if (op == "expm1") {
return t[0] * a.Exp();
}
if (op == "log1p") {
return t[0] / (Array::Ones(a.shape(), a.Dtype()) + a);
}
if (op == "tan") {
Array tn = a.Tan();
return t[0] * (Array::Ones(a.shape(), a.Dtype()) + tn * tn);
}
if (op == "pow" && p.size() == 2) {
Array one = Array::Full(p[1].shape(), 1.0, p[1].Dtype());
return t[0] * p[1] * a.Pow(p[1] - one) + t[1] * a.Pow(p[1]) * a.Log();
}
if (op == "maximum" && p.size() == 2) {
return t[0] * a.GreaterEqual(p[1]).Astype(a.Dtype()) + t[1] * p[1].GreaterThan(a).Astype(p[1].Dtype());
}
if (op == "minimum" && p.size() == 2) {
return t[0] * a.LessEqual(p[1]).Astype(a.Dtype()) + t[1] * p[1].LessThan(a).Astype(p[1].Dtype());
}
if (op == "select" && p.size() == 3) {
return plax::Select(p[0], t[1], t[2]);
}
TENZR_THROW("tangent rule not implemented for custom op: " + std::string(op));
}
} // namespace
bool IsZeroDerivativeCustomOp(std::string_view op) {
return op.rfind("compare_", 0) == 0 || op.rfind("bitwise_", 0) == 0 || op.rfind("shift_", 0) == 0 || op == "logical_not" || op == "is_finite" ||
op == "sign" || op == "ceil" || op == "floor" || op == "round";
}
bool HasDerivativeRule(std::string_view op) {
return op == "tanh" || op == "neg" || op == "exp" || op == "log" || op == "sqrt" || op == "sin" || op == "cos" || op == "abs" || op == "rsqrt" ||
op == "cbrt" || op == "expm1" || op == "log1p" || op == "tan" || op == "pow" || op == "maximum" || op == "minimum" || op == "select" ||
op == "atan2" || op == "remainder" || op == "clamp";
}
bool HasJvpRule(std::string_view op) {
return op == "tanh" || op == "neg" || op == "exp" || op == "log" || op == "sqrt" || op == "sin" || op == "cos" || op == "abs" || op == "rsqrt" ||
op == "cbrt" || op == "expm1" || op == "log1p" || op == "tan" || op == "pow" || op == "maximum" || op == "minimum" || op == "select";
}
void CustomOpVjp(std::string_view op, const std::vector<Array>& inputs, const Array& grad_out, std::vector<Array>& out) { vjp(op, inputs, grad_out, out); }
Array CustomOpJvp(std::string_view op, const std::vector<Array>& primals, const std::vector<Array>& tangents) { return jvp(op, primals, tangents); }
} // namespace tenzr::txpr