// 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 "bindings.h"
namespace tenzr::python {
void BindPlax(nb::module_& m) {
auto plax_m = m.def_submodule("plax", "Low-level strict primitives for OpenXLA (lax-equivalent)");
// Binary Math
plax_m.def("add", &plax::Add, "lhs"_a, "rhs"_a);
plax_m.def("sub", &plax::Sub, "lhs"_a, "rhs"_a);
plax_m.def("mul", &plax::Mul, "lhs"_a, "rhs"_a);
plax_m.def("div", &plax::Div, "lhs"_a, "rhs"_a);
plax_m.def("pow", &plax::Pow, "lhs"_a, "rhs"_a);
plax_m.def("maximum", &plax::Maximum, "lhs"_a, "rhs"_a);
plax_m.def("minimum", &plax::Minimum, "lhs"_a, "rhs"_a);
plax_m.def("remainder", &plax::Remainder, "lhs"_a, "rhs"_a);
plax_m.def("atan2", &plax::Atan2, "lhs"_a, "rhs"_a);
// Comparisons
plax_m.def("eq", &plax::Eq, "lhs"_a, "rhs"_a);
plax_m.def("ne", &plax::Ne, "lhs"_a, "rhs"_a);
plax_m.def("lt", &plax::Lt, "lhs"_a, "rhs"_a);
plax_m.def("le", &plax::Le, "lhs"_a, "rhs"_a);
plax_m.def("gt", &plax::Gt, "lhs"_a, "rhs"_a);
plax_m.def("ge", &plax::Ge, "lhs"_a, "rhs"_a);
// Bitwise & shifts
plax_m.def("bitwise_and", [](const Array& a, const Array& b) { return a.BitwiseAnd(b); }, "lhs"_a, "rhs"_a);
plax_m.def("bitwise_or", [](const Array& a, const Array& b) { return a.BitwiseOr(b); }, "lhs"_a, "rhs"_a);
plax_m.def("bitwise_xor", [](const Array& a, const Array& b) { return a.BitwiseXor(b); }, "lhs"_a, "rhs"_a);
plax_m.def("shift_left", [](const Array& a, const Array& b) { return a.ShiftLeft(b); }, "lhs"_a, "rhs"_a);
plax_m.def("shift_right_logical", [](const Array& a, const Array& b) { return a.ShiftRight(b); }, "lhs"_a, "rhs"_a);
plax_m.def("shift_right_arithmetic", [](const Array& a, const Array& b) { return a.ShiftRight(b); }, "lhs"_a, "rhs"_a);
// Unary Math
plax_m.def("neg", &plax::Neg, "x"_a);
plax_m.def("exp", &plax::Exp, "x"_a);
plax_m.def("log", &plax::Log, "x"_a);
plax_m.def("tanh", &plax::tanh, "x"_a);
plax_m.def("abs", &plax::Abs, "x"_a);
plax_m.def("sqrt", &plax::Sqrt, "x"_a);
plax_m.def("rsqrt", &plax::Rsqrt, "x"_a);
plax_m.def("sin", &plax::Sin, "x"_a);
plax_m.def("cos", &plax::Cos, "x"_a);
plax_m.def("tan", &plax::Tan, "x"_a);
plax_m.def("ceil", &plax::Ceil, "x"_a);
plax_m.def("floor", &plax::Floor, "x"_a);
plax_m.def("round", &plax::Round, "x"_a);
plax_m.def("sign", &plax::Sign, "x"_a);
plax_m.def("cbrt", [](const Array& x) { return x.Cbrt(); }, "x"_a);
plax_m.def("expm1", [](const Array& x) { return x.Expm1(); }, "x"_a);
plax_m.def("log1p", [](const Array& x) { return x.Log1p(); }, "x"_a);
plax_m.def("is_finite", [](const Array& x) { return x.IsFinite(); }, "x"_a);
plax_m.def("logical_not", [](const Array& x) { return x.LogicalNot(); }, "x"_a);
plax_m.def("square", [](const Array& x) { return x * x; }, "x"_a);
plax_m.def("reciprocal", [](const Array& x) { return Array::Full(x.shape(), 1.0, x.Dtype()) / x; }, "x"_a);
// Reductions
plax_m.def("reduce_sum", &plax::ReduceSum, "operand"_a, "axes"_a = std::vector<int64_t>{}, "keepdims"_a = false);
plax_m.def("reduce_max", &plax::ReduceMax, "operand"_a, "axes"_a = std::vector<int64_t>{}, "keepdims"_a = false);
plax_m.def("reduce_min", &plax::ReduceMin, "operand"_a, "axes"_a = std::vector<int64_t>{}, "keepdims"_a = false);
plax_m.def("reduce_prod", &plax::ReduceProd, "operand"_a, "axes"_a = std::vector<int64_t>{}, "keepdims"_a = false);
plax_m.def("reduce_window", &plax::ReduceWindow, "operand"_a, "init_value"_a, "window_dimensions"_a, "window_strides"_a = std::vector<int64_t>{},
"padding_low"_a = std::vector<int64_t>{}, "padding_high"_a = std::vector<int64_t>{}, "reducer"_a = "add");
// Array manipulation & indexing
plax_m.def("slice", &plax::slice, "operand"_a, "start_indices"_a, "limit_indices"_a, "strides"_a = std::vector<int64_t>{});
plax_m.def("pad", &plax::Pad, "operand"_a, "padding_value"_a, "edge_padding_low"_a, "edge_padding_high"_a, "interior_padding"_a);
plax_m.def("gather", &plax::Gather, "operand"_a, "start_indices"_a, "axis"_a = 0);
plax_m.def("dynamic_slice", static_cast<Array (*)(const Array&, const Array&, const std::vector<int64_t>&)>(&plax::DynamicSlice), "operand"_a,
"start_indices"_a, "slice_sizes"_a);
plax_m.def("dynamic_update_slice", static_cast<Array (*)(const Array&, const Array&, const Array&)>(&plax::DynamicUpdateSlice), "operand"_a, "update"_a,
"start_indices"_a);
plax_m.def("scatter_add", &plax::ScatterAdd, "operand"_a, "indices"_a, "updates"_a, "axis"_a = 0);
plax_m.def("scatter", &plax::Scatter, "operand"_a, "indices"_a, "updates"_a, "axis"_a = 0, "update"_a = "add");
plax_m.def("argsort", &plax::Argsort, "operand"_a, "axis"_a = -1, "descending"_a = false);
plax_m.def("sort", &plax::Sort, "operand"_a, "axis"_a = -1, "descending"_a = false);
plax_m.def("argmax", &plax::Argmax, "operand"_a, "axis"_a = -1);
plax_m.def("argmin", &plax::Argmin, "operand"_a, "axis"_a = -1);
plax_m.def("concatenate", &plax::Concatenate, "arrays"_a, "axis"_a = 0);
plax_m.def("stack", &plax::Stack, "arrays"_a, "axis"_a = 0);
plax_m.def("conv2d", &plax::Conv2d, "input"_a, "weight"_a, "stride"_a = std::pair<int64_t, int64_t>{1, 1}, "padding"_a = std::pair<int64_t, int64_t>{0, 0});
plax_m.def("take_along_axis", &plax::TakeAlongAxis, "operand"_a, "indices"_a, "axis"_a = -1);
plax_m.def("put_along_axis", &plax::PutAlongAxis, "operand"_a, "indices"_a, "values"_a, "axis"_a = -1, "update"_a = "set");
plax_m.def(
"top_k",
[](const Array& operand, int64_t k, int64_t axis, bool largest) -> nb::tuple {
auto r = plax::TopK(operand, k, axis, largest);
return nb::make_tuple(r.values, r.indices);
},
"operand"_a, "k"_a, "axis"_a = -1, "largest"_a = true);
plax_m.def("reverse", &plax::Reverse, "operand"_a, "dimensions"_a);
// Dot & Transforms
plax_m.def("dot_general", &plax::DotGeneral, "lhs"_a, "rhs"_a, "lhs_contracting_dims"_a, "rhs_contracting_dims"_a,
"lhs_batch_dims"_a = std::vector<int64_t>{}, "rhs_batch_dims"_a = std::vector<int64_t>{});
plax_m.def(
"broadcast_in_dim",
[](const Array& op, std::vector<int64_t> target_shape, std::vector<int64_t> bcast_dims) {
return plax::BroadcastInDim(op, Shape(target_shape), bcast_dims);
},
"operand"_a, "target_shape"_a, "broadcast_dimensions"_a);
plax_m.def("reshape", [](const Array& op, std::vector<int64_t> new_shape) { return plax::Reshape(op, Shape(new_shape)); }, "operand"_a, "new_shape"_a);
plax_m.def("transpose", &plax::Transpose, "operand"_a, "permutation"_a);
plax_m.def("convert_element_type", &plax::ConvertElementType, "operand"_a, "new_dtype"_a);
plax_m.def("select", &plax::Select, "pred"_a, "on_true"_a, "on_false"_a);
plax_m.def("clamp", &plax::Clamp, "min_val"_a, "operand"_a, "max_val"_a);
// Distributed Collectives
plax_m.def("all_reduce", &plax::AllReduce, "x"_a, "op"_a = "sum", "replica_groups"_a = std::vector<std::vector<int64_t>>{});
plax_m.def("all_gather", &plax::AllGather, "x"_a, "all_gather_dim"_a = 0, "replica_groups"_a = std::vector<std::vector<int64_t>>{});
plax_m.def("all_to_all", &plax::AllToAll, "x"_a, "split_dim"_a = 0, "concat_dim"_a = 0, "split_count"_a = 1,
"replica_groups"_a = std::vector<std::vector<int64_t>>{});
plax_m.def("reduce_scatter", &plax::ReduceScatter, "x"_a, "op"_a = "sum", "scatter_dim"_a = 0, "replica_groups"_a = std::vector<std::vector<int64_t>>{});
plax_m.def("psum", &plax::Psum, "x"_a, "axis_name"_a = "");
plax_m.def("pmean", &plax::Pmean, "x"_a, "axis_name"_a = "");
plax_m.def("with_sharding_constraint", &plax::WithShardingConstraint, "x"_a, "sharding"_a);
// Structured Control Flow
plax_m.def(
"cond",
[](const Array& pred, nb::callable true_fun, nb::callable false_fun, const Array& operand) -> Array {
auto cpp_true = [true_fun](const Array& a) -> Array { return nb::cast<Array>(true_fun(a)); };
auto cpp_false = [false_fun](const Array& a) -> Array { return nb::cast<Array>(false_fun(a)); };
return txpr::Cond(pred, cpp_true, cpp_false, operand);
},
"pred"_a, "true_fun"_a, "false_fun"_a, "operand"_a);
plax_m.def(
"scan",
[](nb::callable f, const Array& init, const Array& xs, int64_t length) -> nb::tuple {
auto cpp_f = [f](const Array& c, const Array& x) -> txpr::ScanResult {
nb::tuple res = nb::cast<nb::tuple>(f(c, x));
return {nb::cast<Array>(res[0]), nb::cast<Array>(res[1])};
};
auto r = txpr::Scan(cpp_f, init, xs, length);
return nb::make_tuple(r.carry, r.ys);
},
"f"_a, "init"_a, "xs"_a, "length"_a);
plax_m.def(
"while_loop",
[](nb::callable cond_fun, nb::callable body_fun, const Array& init_val) -> Array {
auto cpp_cond = [cond_fun](const Array& a) -> Array { return nb::cast<Array>(cond_fun(a)); };
auto cpp_body = [body_fun](const Array& a) -> Array { return nb::cast<Array>(body_fun(a)); };
return txpr::WhileLoop(cpp_cond, cpp_body, init_val);
},
"cond_fun"_a, "body_fun"_a, "init_val"_a);
}
} // namespace tenzr::python