// 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 <memory>
#include "bindings.h"
namespace tenzr::python {
namespace {
std::vector<int> ParseInAxes(nb::handle in_axes) {
std::vector<int> axes;
if (in_axes.is_none()) {
return axes;
}
if (nb::isinstance<nb::int_>(in_axes)) {
axes.push_back(nb::cast<int>(in_axes));
return axes;
}
if (nb::isinstance<nb::tuple>(in_axes) || nb::isinstance<nb::list>(in_axes)) {
for (auto item : nb::cast<nb::sequence>(in_axes)) {
axes.push_back(item.is_none() ? 0 : nb::cast<int>(item));
}
return axes;
}
throw nb::type_error("in_axes must be an int, a sequence of ints, or None");
}
} // namespace
void BindTransforms(nb::module_& m) {
m.def(
"grad",
[](nb::callable fn, int64_t argnums) {
if (argnums != 0) {
throw nb::type_error(
"tenzr.grad currently differentiates only argument 0; "
"curry extra arguments into the loss function (argnums != 0 is not yet supported)");
}
auto cpp_fn = [fn](const Array& a) -> Array { return nb::cast<Array>(fn(a)); };
auto g_fn = std::make_shared<std::function<Array(const Array&)>>(tenzr::Grad(cpp_fn));
return nb::cpp_function([g_fn](const Array& x) -> Array { return (*g_fn)(x); });
},
"fn"_a, "argnums"_a = 0);
m.def(
"value_and_grad",
[](nb::callable fn, int64_t argnums) {
if (argnums != 0) {
throw nb::type_error(
"tenzr.value_and_grad currently differentiates only argument 0; "
"curry extra arguments into the loss function (argnums != 0 is not yet supported)");
}
auto cpp_fn = [fn](const Array& a) -> Array { return nb::cast<Array>(fn(a)); };
auto vg_fn = std::make_shared<std::function<txpr::ValueAndGradResult(const Array&)>>(tenzr::ValueAndGrad(cpp_fn));
return nb::cpp_function([vg_fn](const Array& x) -> nb::tuple {
auto r = (*vg_fn)(x);
return nb::make_tuple(r.value, r.grad);
});
},
"fn"_a, "argnums"_a = 0);
// Multi-argument jit (arity 1-3). The core JittedFunction is variadic.
m.def(
"jit",
[](nb::callable fn) {
auto unary = [fn](const Array& a) -> Array { return nb::cast<Array>(fn(a)); };
auto binary = [fn](const Array& a, const Array& b) -> Array { return nb::cast<Array>(fn(a, b)); };
auto ternary = [fn](const Array& a, const Array& b, const Array& c) -> Array { return nb::cast<Array>(fn(a, b, c)); };
auto j1 = std::make_shared<decltype(tenzr::Jit(unary))>(tenzr::Jit(unary));
auto j2 = std::make_shared<decltype(tenzr::Jit(binary))>(tenzr::Jit(binary));
auto j3 = std::make_shared<decltype(tenzr::Jit(ternary))>(tenzr::Jit(ternary));
return nb::cpp_function([j1, j2, j3](nb::args args) -> Array {
switch (args.size()) {
case 1:
return (*j1)(nb::cast<Array>(args[0]));
case 2:
return (*j2)(nb::cast<Array>(args[0]), nb::cast<Array>(args[1]));
case 3:
return (*j3)(nb::cast<Array>(args[0]), nb::cast<Array>(args[1]), nb::cast<Array>(args[2]));
default:
throw nb::type_error("tenzr.jit supports between 1 and 3 positional array arguments");
}
});
},
"fn"_a);
// Multi-argument vmap (arity 1-3) with int or per-argument in_axes.
m.def(
"vmap",
[](nb::callable fn, nb::object in_axes, int64_t out_axis) {
std::vector<int> axes = ParseInAxes(in_axes);
auto unary = [fn](const Array& a) -> Array { return nb::cast<Array>(fn(a)); };
auto binary = [fn](const Array& a, const Array& b) -> Array { return nb::cast<Array>(fn(a, b)); };
auto ternary = [fn](const Array& a, const Array& b, const Array& c) -> Array { return nb::cast<Array>(fn(a, b, c)); };
auto v1 = std::make_shared<decltype(tenzr::Vmap(unary, axes, static_cast<int>(out_axis)))>(tenzr::Vmap(unary, axes, static_cast<int>(out_axis)));
auto v2 = std::make_shared<decltype(tenzr::Vmap(binary, axes, static_cast<int>(out_axis)))>(tenzr::Vmap(binary, axes, static_cast<int>(out_axis)));
auto v3 = std::make_shared<decltype(tenzr::Vmap(ternary, axes, static_cast<int>(out_axis)))>(tenzr::Vmap(ternary, axes, static_cast<int>(out_axis)));
return nb::cpp_function([v1, v2, v3](nb::args args) -> Array {
switch (args.size()) {
case 1:
return (*v1)(nb::cast<Array>(args[0]));
case 2:
return (*v2)(nb::cast<Array>(args[0]), nb::cast<Array>(args[1]));
case 3:
return (*v3)(nb::cast<Array>(args[0]), nb::cast<Array>(args[1]), nb::cast<Array>(args[2]));
default:
throw nb::type_error("tenzr.vmap supports between 1 and 3 positional array arguments");
}
});
},
"fn"_a, "in_axes"_a = 0, "out_axes"_a = 0);
m.def(
"jvp",
[](nb::callable fn, const Array& primal, const Array& tangent) -> nb::tuple {
auto cpp_fn = [fn](const Array& a) -> Array { return nb::cast<Array>(fn(a)); };
auto r = txpr::Jvp(cpp_fn, primal, tangent);
return nb::make_tuple(r.primal, r.tangent);
},
"fn"_a, "primals"_a, "tangents"_a);
m.def(
"linearize",
[](nb::callable fn, const Array& primal) {
auto cpp_fn = [fn](const Array& a) -> Array { return nb::cast<Array>(fn(a)); };
auto [y, jvp_fn] = txpr::Linearize(cpp_fn, primal);
auto py_jvp = nb::cpp_function([jvp_fn](const Array& tangent) -> Array { return jvp_fn(tangent); });
return nb::make_tuple(y, py_jvp);
},
"fn"_a, "primals"_a);
m.def(
"vjp",
[](nb::callable fn, nb::sequence primals) {
std::vector<Array> ps;
for (auto item : primals) {
ps.push_back(nb::cast<Array>(item));
}
auto cpp_fn = [fn](const std::vector<Array>& xs) -> Array {
nb::list args;
for (const auto& x : xs) {
args.append(x);
}
return nb::cast<Array>(fn(args));
};
auto result = txpr::Vjp(cpp_fn, ps);
auto pullback = result.pullback;
auto py_pullback = nb::cpp_function([pullback](const Array& cotangent) {
nb::list out;
for (const auto& g : pullback(cotangent)) {
out.append(g);
}
return out;
});
return nb::make_tuple(result.output, py_pullback);
},
"fn"_a, "primals"_a);
}
} // namespace tenzr::python