// 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.
#ifndef TENZR_TXPR_VMAP_H_
#define TENZR_TXPR_VMAP_H_
#include <cstdint>
#include <functional>
#include <memory>
#include <tuple>
#include <type_traits>
#include <unordered_map>
#include <utility>
#include <vector>
#include "tenzr/common/error.h"
#include "tenzr/execution/jit_cache.h"
#include "tenzr/frontend/array.h"
#include "tenzr/txpr/emitter.h"
#include "tenzr/txpr/jit.h"
#include "tenzr/txpr/tracer.h"
#include "tenzr/txpr/txpr.h"
namespace tenzr::txpr {
Txpr BatchTxpr(const Txpr& orig, int64_t batch_size, const std::vector<int>& in_axes = {0}, int out_axis = 0);
template <typename Fn> class VmappedFunction {
public:
explicit VmappedFunction(Fn fn, std::vector<int> in_axes = {0}, int out_axis = 0)
: fn_(std::move(fn)), in_axes_(std::move(in_axes)), out_axis_(out_axis), fn_token_(detail::NextJitToken()) {}
template <typename... Args> Array operator()(const Args&... args) {
std::vector<Array> in_args = {args...};
TENZR_CHECK(!in_args.empty(), "vmap requires at least one input argument");
std::vector<int> effective_in_axes = in_axes_;
while (effective_in_axes.size() < in_args.size()) {
effective_in_axes.push_back(0);
}
int64_t batch_size = -1;
for (size_t i = 0; i < in_args.size(); ++i) {
if (effective_in_axes[i] >= 0) {
int ax = effective_in_axes[i];
TENZR_CHECK(in_args[i].shape().Rank() > static_cast<size_t>(ax), "vmap: input rank must exceed in_axis");
int64_t b = in_args[i].shape().Dim(static_cast<size_t>(ax));
if (batch_size < 0) {
batch_size = b;
} else {
TENZR_CHECK(batch_size == b, "vmap: inconsistent batch dimensions across inputs");
}
}
}
TENZR_CHECK(batch_size > 0, "vmap: could not determine positive batch size");
// 1. Trace the kernel unbatched
Txpr unbatched_txpr;
std::vector<Array> sym_unbatched_args;
sym_unbatched_args.reserve(in_args.size());
for (size_t i = 0; i < in_args.size(); ++i) {
Shape unbatched_shape;
if (effective_in_axes[i] >= 0) {
std::vector<int64_t> dims = in_args[i].shape().Dims();
dims.erase(dims.begin() + effective_in_axes[i]);
unbatched_shape = Shape(dims);
} else {
unbatched_shape = in_args[i].shape();
}
VarId vid = unbatched_txpr.AddVar(unbatched_shape, in_args[i].Dtype());
unbatched_txpr.MarkInput(vid);
sym_unbatched_args.push_back(Array::Symbolic(vid, unbatched_shape, in_args[i].Dtype()));
}
auto sym_tuple = [&sym_unbatched_args]<size_t... Is>(std::index_sequence<Is...>) {
return std::forward_as_tuple(sym_unbatched_args[Is]...);
}(std::make_index_sequence<sizeof...(Args)>{});
Array sym_out;
{
TracerScope scope(unbatched_txpr);
sym_out = std::apply(fn_, sym_tuple);
}
TENZR_CHECK(sym_out.IsSymbolic(), "vmap: function output must be symbolic during tracing");
unbatched_txpr.MarkOutput(sym_out.SymbolicVarId());
// 2. Vectorize the unbatched Txpr
Txpr batched_txpr = BatchTxpr(unbatched_txpr, batch_size, effective_in_axes, out_axis_);
// 3. If an outer tracer is active (e.g. inside tenzr::jit), graft batched equations
if (TracerContext::IsActive()) {
Txpr* outer_txpr = TracerContext::CurrentTxpr();
std::unordered_map<uint32_t, VarId> vmap_to_outer;
for (size_t i = 0; i < in_args.size(); ++i) {
vmap_to_outer[batched_txpr.InBinders()[i].id] = in_args[i].SymbolicVarId();
}
for (const auto& eqn : batched_txpr.Equations()) {
std::vector<VarId> outer_ins;
outer_ins.reserve(eqn.in_vars.size());
for (const auto& in_v : eqn.in_vars) {
outer_ins.push_back(vmap_to_outer.at(in_v.id));
}
const auto& out_def = batched_txpr.Variable(eqn.out_vars[0]);
VarId outer_out = outer_txpr->AppendInstruction(eqn.op, outer_ins, out_def.shape, out_def.dtype, eqn.attrs);
vmap_to_outer[eqn.out_vars[0].id] = outer_out;
}
VarId outer_res_id = vmap_to_outer.at(batched_txpr.OutBinders()[0].id);
const auto& res_def = batched_txpr.Variable(batched_txpr.OutBinders()[0]);
return Array::Symbolic(outer_res_id, res_def.shape, res_def.dtype);
}
// 4. Standalone direct dispatch
uint64_t hash = detail::HashSignature(fn_token_, in_args);
auto cached = JitCache::Instance().Get(JitCache::DomainKey(JitCache::kJitKeyDomain, hash));
if (!cached) {
cached = CompileTxprModule(hash, batched_txpr);
}
ExecutionResult res = ExecuteJittedExecutable(cached, in_args);
return res.AsArray();
}
private:
Fn fn_;
std::vector<int> in_axes_;
int out_axis_{0};
uint64_t fn_token_;
};
template <typename Fn>
requires(!std::is_same_v<std::decay_t<Fn>, Array>)
auto Vmap(Fn&& fn, std::vector<int> in_axes = {0}, int out_axis = 0) {
return VmappedFunction<std::decay_t<Fn>>(std::forward<Fn>(fn), std::move(in_axes), out_axis);
}
template <typename Fn>
requires(!std::is_same_v<std::decay_t<Fn>, Array>)
auto Vmap(Fn&& fn, int in_axis, int out_axis = 0) {
return VmappedFunction<std::decay_t<Fn>>(std::forward<Fn>(fn), std::vector<int>{in_axis}, out_axis);
}
} // namespace tenzr::txpr
namespace tenzr {
using txpr::BatchTxpr;
using txpr::Vmap;
using txpr::VmappedFunction;
} // namespace tenzr
#endif // TENZR_TXPR_VMAP_H_