tiny.accy.tensor.batch

1const std = @import("std");2const emit = @import("emit.zig");3const hook = @import("hook.zig");4const interpret = @import("interpret/root.zig");5const program_mod = @import("program.zig");6const trace = @import("trace/root.zig");7const transform = @import("transform.zig");8const types = @import("type/root.zig");910pub const Axis = union(enum) {11    none,12    axis: usize,13};1415pub const Options = struct {16    axis_size: i64,17    in_axes: []const Axis,18    out_axis: usize = 0,19    axis_name: []const u8 = &.{},20};2122const default_axis_name = "#batch";2324const BatchValue = struct {25    value: trace.Value,26    axis: ?usize,27};2829pub fn mapped(axis_index: usize) Axis {30    return .{ .axis = axis_index };31}3233pub fn vmap(34    allocator: std.mem.Allocator,35    source: *const program_mod.Program,36    options: Options,37) !program_mod.Program {38    try validate(source, options);3940    var builder = try trace.Builder.init(allocator, source.name);41    errdefer builder.deinit();4243    const resolved = try resolveAxisName(builder.arena.allocator(), source, options);44    const graph = interpret.Graph{ .builder = &builder };45    return try interpret.run(allocator, source, semantics(graph, resolved));46}4748pub fn vmapWith(49    allocator: std.mem.Allocator,50    source: *const program_mod.Program,51    options: Options,52    hooks: anytype,53) !program_mod.Program {54    try validate(source, options);5556    var builder = try trace.Builder.init(allocator, source.name);57    errdefer builder.deinit();5859    const resolved = try resolveAxisName(builder.arena.allocator(), source, options);60    const graph = interpret.Graph{ .builder = &builder };61    return try interpret.run(allocator, source, semantics(hook.attach("batch", hooks, graph), resolved));62}6364fn resolveAxisName(allocator: std.mem.Allocator, source: *const program_mod.Program, options: Options) !Options {65    var resolved = options;66    if (options.axis_name.len != 0) {67        try types.validateName(options.axis_name);68        return resolved;69    }70    var attempt: usize = 0;71    var name: []const u8 = default_axis_name;72    while (programHasAxis(source, name)) {73        attempt += 1;74        name = try std.fmt.allocPrint(allocator, "{s}{d}", .{ default_axis_name, attempt });75    }76    resolved.axis_name = name;77    return resolved;78}7980fn programHasAxis(source: *const program_mod.Program, name: []const u8) bool {81    for (source.values) |ty| {82        if (types.findDim(ty.dims, name) != null) return true;83    }84    for (source.operations) |op| {85        if (op.kind == .scan and subgraphHasAxis(op.kind.scan.body, name)) return true;86    }87    return false;88}8990fn subgraphHasAxis(body: *const program_mod.Subgraph, name: []const u8) bool {91    for (body.values) |ty| {92        if (types.findDim(ty.dims, name) != null) return true;93    }94    for (body.operations) |op| {95        if (op.kind == .scan and subgraphHasAxis(op.kind.scan.body, name)) return true;96    }97    return false;98}99100pub fn semantics(next: anytype, options: Options) Semantics(@TypeOf(next)) {101    return .{102        .next = next,103        .options = options,104    };105}106107pub fn semanticsWith(next: anytype, options: Options, hooks: anytype) Semantics(@TypeOf(hook.attach("batch", hooks, next))) {108    return semantics(hook.attach("batch", hooks, next), options);109}110111pub fn Semantics(comptime Next: type) type {112    return struct {113        next: Next,114        options: Options,115116        pub const Value: type = BatchValue;117        pub const Result: type = Next.Result;118119        pub fn operation(self: *@This(), step: *interpret.Step(Value)) !Value {120            return batchOperation(self, step.op, step.values);121        }122123        pub fn finish(self: *@This(), outputs: []const Value) !Result {124            const values = try self.builderHandle().arena.allocator().alloc(trace.Value, outputs.len);125            for (outputs, 0..) |output, index| {126                values[index] = if (output.axis) |axis|127                    try moveBatchAxis(self, output.value, axis, self.options.out_axis)128                else129                    try broadcastOutputAxis(self, output.value, self.options);130            }131            return self.next.finish(values);132        }133134        pub fn builderHandle(self: *@This()) *trace.Builder {135            return self.next.builderHandle();136        }137138        pub fn axisName(self: *@This()) []const u8 {139            if (self.options.axis_name.len != 0) return self.options.axis_name;140            return default_axis_name;141        }142    };143}144145fn batchDim(layer: anytype) types.Dim {146    return .{ .name = layer.axisName(), .extent = layer.options.axis_size };147}148149fn batchOperation(layer: anytype, op: *const program_mod.Operation, map: []const BatchValue) !BatchValue {150    return switch (op.kind) {151        .parameter => |parameter| batchParameter(layer, op, parameter.index),152        .constant, .iota => .{ .value = try emitUnbatched(layer, op, map), .axis = null },153        .unary => |unary| blk: {154            const input = map[unary.input.index];155            break :blk .{ .value = try layer.next.bind(op, &.{input.value}), .axis = input.axis };156        },157        .binary => |binary| batchBinary(layer, op, map[binary.lhs.index], map[binary.rhs.index]),158        .reshape => |reshape| blk: {159            const input = map[reshape.input.index];160            if (input.axis) |axis| {161                const allocator = layer.builderHandle().arena.allocator();162                const dims = try types.insertDim(allocator, op.result.dims, axis, batchDim(layer));163                var batched_op = op.*;164                batched_op.result = .{ .dtype = op.result.dtype, .dims = dims };165                batched_op.kind = .{ .reshape = .{ .input = reshape.input, .new_shape = try types.extents(allocator, dims) } };166                break :blk .{ .value = try layer.next.bind(&batched_op, &.{input.value}), .axis = axis };167            }168            break :blk .{ .value = try emitUnbatched(layer, op, map), .axis = null };169        },170        .broadcast => |broadcast| batchBroadcast(layer, op, broadcast, map[broadcast.input.index]),171        .broadcast_in_dim => |broadcast| batchBroadcastInDim(layer, op, broadcast, map[broadcast.input.index]),172        .transpose => |transpose| batchTranspose(layer, op, transpose, map[transpose.input.index]),173        .reduce => |reduce| batchReduce(layer, op, reduce, map[reduce.input.index], map[reduce.init.index]),174        .gather => |gather| batchGather(layer, op, map[gather.input.index], map[gather.indices.index]),175        .scatter_add => |scatter_add| batchScatterAdd(176            layer,177            op,178            map[scatter_add.input.index],179            map[scatter_add.indices.index],180            map[scatter_add.updates.index],181        ),182        .sparse_cross_entropy => |sparse_cross_entropy| batchSparseCrossEntropy(183            layer,184            op,185            map[sparse_cross_entropy.logits.index],186            map[sparse_cross_entropy.targets.index],187        ),188        .dot_general => |dot| batchDot(layer, op, dot, map[dot.lhs.index], map[dot.rhs.index]),189        .compare => |compare| batchBinary(layer, op, map[compare.lhs.index], map[compare.rhs.index]),190        .select => |select| batchSelect(layer, op, map[select.pred.index], map[select.on_true.index], map[select.on_false.index]),191        .custom_call => error.CustomCallRequiresBatchContract,192        .scan => |scan| batchScan(layer, op, scan, map),193        .projection => |projection| blk: {194            const source = map[projection.source.index];195            break :blk .{ .value = try layer.next.bind(op, &.{source.value}), .axis = source.axis };196        },197    };198}199200fn batchScan(layer: anytype, op: *const program_mod.Operation, scan: program_mod.Scan, map: []const BatchValue) anyerror!BatchValue {201    var any_batched = false;202    for (scan.inits) |init_id| {203        if (map[init_id.index].axis != null) any_batched = true;204    }205    if (!any_batched) {206        var buffer: [program_mod.max_scan_carries]trace.Value = undefined;207        for (scan.inits, 0..) |init_id, index| buffer[index] = map[init_id.index].value;208        return .{ .value = try layer.next.bind(op, buffer[0..scan.inits.len]), .axis = null };209    }210211    const builder = layer.builderHandle();212    var inits: [program_mod.max_scan_carries]trace.Value = undefined;213    for (scan.inits, 0..) |init_id, index| {214        const carried = map[init_id.index];215        inits[index] = if (carried.axis) |axis|216            try moveBatchAxis(layer, carried.value, axis, 0)217        else218            try broadcastToBatchAxis(layer, carried.value, 0);219    }220221    var body_program = program_mod.Program{222        .arena = std.heap.ArenaAllocator.init(builder.allocator),223        .name = "scan_body",224        .values = scan.body.values,225        .operations = scan.body.operations,226        .parameters = scan.body.parameters,227        .outputs = scan.body.outputs,228    };229    defer body_program.deinit();230231    var axes: [program_mod.max_scan_carries]Axis = undefined;232    for (0..scan.inits.len) |index| axes[index] = mapped(0);233    var batched_body = try vmap(builder.allocator, &body_program, .{234        .axis_size = layer.options.axis_size,235        .in_axes = axes[0..scan.inits.len],236        .axis_name = layer.axisName(),237    });238    defer batched_body.deinit();239240    const body_view = program_mod.Subgraph{241        .values = batched_body.values,242        .operations = batched_body.operations,243        .parameters = batched_body.parameters,244        .outputs = batched_body.outputs,245    };246    var batched_op = program_mod.Operation{247        .id = program_mod.synthetic_id,248        .result = inits[0].ty,249        .kind = .{ .scan = .{ .length = scan.length, .inits = scan.inits, .body = &body_view } },250    };251    return .{ .value = try layer.next.bind(&batched_op, inits[0..scan.inits.len]), .axis = 0 };252}253254fn batchGather(layer: anytype, op: *const program_mod.Operation, input: BatchValue, indices: BatchValue) !BatchValue {255    if (input.axis == null and indices.axis == null) {256        return .{ .value = try layer.next.bind(op, &.{ input.value, indices.value }), .axis = null };257    }258    const gather = op.kind.gather;259    if (input.axis != null and indices.axis != null) {260        const moved_input = try moveBatchAxis(layer, input.value, input.axis.?, 0);261        const moved_indices = try moveBatchAxis(layer, indices.value, indices.axis.?, 0);262        return .{263            .value = try pairedGather(layer, moved_input, moved_indices, try leadingAxis(gather.axis)),264            .axis = 0,265        };266    }267    if (input.axis) |axis| {268        const moved_input = try moveBatchAxis(layer, input.value, axis, 0);269        return .{270            .value = try emit.gather(layer, moved_input, indices.value, try leadingAxis(gather.axis)),271            .axis = 0,272        };273    }274    const moved_indices = try moveBatchAxis(layer, indices.value, indices.axis.?, 0);275    return .{276        .value = try emit.gather(layer, input.value, moved_indices, gather.axis),277        .axis = @intCast(gather.axis),278    };279}280281fn batchScatterAdd(layer: anytype, op: *const program_mod.Operation, input: BatchValue, indices: BatchValue, updates: BatchValue) !BatchValue {282    if (input.axis == null and indices.axis == null and updates.axis == null) {283        return .{ .value = try layer.next.bind(op, &.{ input.value, indices.value, updates.value }), .axis = null };284    }285    const scatter_add = op.kind.scatter_add;286    const moved_input = if (input.axis) |axis|287        try moveBatchAxis(layer, input.value, axis, 0)288    else289        try broadcastToBatchAxis(layer, input.value, 0);290    const moved_updates = if (updates.axis) |axis|291        try moveBatchAxis(layer, updates.value, axis, 0)292    else293        try broadcastToBatchAxis(layer, updates.value, 0);294    if (indices.axis) |axis| {295        const moved_indices = try moveBatchAxis(layer, indices.value, axis, 0);296        return .{297            .value = try pairedScatterAdd(layer, moved_input, moved_indices, moved_updates, try leadingAxis(scatter_add.axis)),298            .axis = 0,299        };300    }301    return .{302        .value = try emit.scatterAdd(layer, moved_input, indices.value, moved_updates, try leadingAxis(scatter_add.axis)),303        .axis = 0,304    };305}306307fn batchSparseCrossEntropy(layer: anytype, op: *const program_mod.Operation, logits: BatchValue, targets: BatchValue) !BatchValue {308    if (logits.axis == null and targets.axis == null) {309        return .{ .value = try layer.next.bind(op, &.{ logits.value, targets.value }), .axis = null };310    }311    const sparse_cross_entropy = op.kind.sparse_cross_entropy;312    const moved_logits = if (logits.axis) |axis|313        try moveBatchAxis(layer, logits.value, axis, 0)314    else315        try broadcastToBatchAxis(layer, logits.value, 0);316    const moved_targets = if (targets.axis) |axis|317        try moveBatchAxis(layer, targets.value, axis, 0)318    else319        try broadcastToBatchAxis(layer, targets.value, 0);320    return .{321        .value = try emit.sparseCrossEntropy(layer, moved_logits, moved_targets, try leadingAxis(sparse_cross_entropy.axis)),322        .axis = 0,323    };324}325326fn pairedGather(layer: anytype, input: trace.Value, indices: trace.Value, axis: i64) !trace.Value {327    if (!input.ty.dtype.isNumeric()) return error.UnsupportedBatching;328    const axis_index = try batchIndexAxis(input, indices, axis);329    const allocator = layer.builderHandle().arena.allocator();330    const index_dims = indices.ty.dims[1..];331    const expanded_dims = try pairedExpandedDims(allocator, input.ty.dims, index_dims, axis_index);332    const source_positions = try emitIota(layer, .i32, expanded_dims, axis_index);333    const broadcasted_indices = try broadcastPairedIndices(layer, indices, expanded_dims, axis_index);334    const broadcasted_input = try broadcastPairedInput(layer, input, expanded_dims, axis_index, index_dims.len);335    const mask = try emit.compare(layer, .eq, source_positions, broadcasted_indices);336    const zero = try broadcastZero(layer, input.ty.dtype, expanded_dims);337    const selected = try emit.select(layer, mask, broadcasted_input, zero);338    const init = try emit.zeros(layer, trace.Type.scalar(input.ty.dtype));339    return emit.reduce(layer, selected, init, .sum, &.{axis});340}341342fn pairedScatterAdd(layer: anytype, input: trace.Value, indices: trace.Value, updates: trace.Value, axis: i64) !trace.Value {343    if (!input.ty.dtype.isNumeric()) return error.UnsupportedBatching;344    const axis_index = try batchIndexAxis(input, indices, axis);345    const allocator = layer.builderHandle().arena.allocator();346    const index_dims = indices.ty.dims[1..];347    const expanded_dims = try pairedExpandedDims(allocator, input.ty.dims, index_dims, axis_index);348    const source_positions = try emitIota(layer, .i32, expanded_dims, axis_index);349    const broadcasted_indices = try broadcastPairedIndices(layer, indices, expanded_dims, axis_index);350    const broadcasted_updates = try broadcastPairedUpdates(layer, updates, expanded_dims, axis_index);351    const mask = try emit.compare(layer, .eq, source_positions, broadcasted_indices);352    const zero = try broadcastZero(layer, input.ty.dtype, expanded_dims);353    const selected = try emit.select(layer, mask, broadcasted_updates, zero);354355    const reduced = if (index_dims.len == 0)356        selected357    else blk: {358        const reduce_axes = try allocator.alloc(i64, index_dims.len);359        for (reduce_axes, 0..) |*slot, index| slot.* = @intCast(axis_index + 1 + index);360        const init = try emit.zeros(layer, trace.Type.scalar(input.ty.dtype));361        break :blk try emit.reduce(layer, selected, init, .sum, reduce_axes);362    };363    return emit.binary(layer, .add, input, reduced);364}365366fn batchIndexAxis(input: trace.Value, indices: trace.Value, axis: i64) !usize {367    if (axis <= 0) return error.AxisOutOfRange;368    const axis_index = std.math.cast(usize, axis) orelse return error.AxisOutOfRange;369    if (axis_index >= input.ty.rank()) return error.AxisOutOfRange;370    if (indices.ty.rank() == 0) return error.AxisOutOfRange;371    if (indices.ty.dims[0].extent != input.ty.dims[0].extent) return error.ShapeMismatch;372    return axis_index;373}374375fn pairedExpandedDims(376    allocator: std.mem.Allocator,377    input_dims: []const trace.Dim,378    index_dims: []const trace.Dim,379    axis: usize,380) ![]const trace.Dim {381    const result = try allocator.alloc(trace.Dim, input_dims.len + index_dims.len);382    var out: usize = 0;383    for (input_dims[0 .. axis + 1]) |dim| {384        result[out] = dim;385        out += 1;386    }387    for (index_dims) |dim| {388        result[out] = dim;389        out += 1;390    }391    for (input_dims[axis + 1 ..]) |dim| {392        result[out] = dim;393        out += 1;394    }395    try types.validateDims(result);396    return result;397}398399fn broadcastPairedInput(layer: anytype, input: trace.Value, expanded_dims: []const trace.Dim, axis: usize, index_rank: usize) !trace.Value {400    const allocator = layer.builderHandle().arena.allocator();401    const mapping = try allocator.alloc(i64, input.ty.rank());402    for (mapping, 0..) |*slot, index| {403        slot.* = if (index <= axis)404            @intCast(index)405        else406            @intCast(index + index_rank);407    }408    return emit.broadcastInDim(layer, input, expanded_dims, mapping);409}410411fn broadcastPairedIndices(layer: anytype, indices: trace.Value, expanded_dims: []const trace.Dim, axis: usize) !trace.Value {412    const allocator = layer.builderHandle().arena.allocator();413    const mapping = try allocator.alloc(i64, indices.ty.rank());414    mapping[0] = 0;415    for (mapping[1..], 0..) |*slot, index| slot.* = @intCast(axis + 1 + index);416    return emit.broadcastInDim(layer, indices, expanded_dims, mapping);417}418419fn broadcastPairedUpdates(layer: anytype, updates: trace.Value, expanded_dims: []const trace.Dim, axis: usize) !trace.Value {420    const allocator = layer.builderHandle().arena.allocator();421    const mapping = try allocator.alloc(i64, updates.ty.rank());422    for (mapping, 0..) |*slot, index| {423        slot.* = if (index < axis)424            @intCast(index)425        else426            @intCast(index + 1);427    }428    return emit.broadcastInDim(layer, updates, expanded_dims, mapping);429}430431fn broadcastZero(layer: anytype, dtype: trace.DType, dims: []const trace.Dim) !trace.Value {432    const zero = try emit.zeros(layer, trace.Type.scalar(dtype));433    return emit.broadcastInDim(layer, zero, dims, &.{});434}435436fn emitIota(layer: anytype, dtype: trace.DType, dims: []const trace.Dim, axis: usize) !trace.Value {437    var op = program_mod.Operation{438        .id = program_mod.synthetic_id,439        .result = .{ .dtype = dtype, .dims = dims },440        .kind = .{ .iota = .{ .axis = @intCast(axis) } },441    };442    return layer.next.bind(&op, &.{});443}444445fn leadingAxis(axis: i64) !i64 {446    if (axis < 0) return error.AxisOutOfRange;447    return std.math.add(i64, axis, 1) catch error.AxisOutOfRange;448}449450fn batchParameter(layer: anytype, op: *const program_mod.Operation, index: usize) !BatchValue {451    return switch (layer.options.in_axes[index]) {452        .none => .{ .value = try layer.next.bind(op, &.{}), .axis = null },453        .axis => |axis| blk: {454            var batched_op = op.*;455            batched_op.result = .{456                .dtype = op.result.dtype,457                .dims = try types.insertDim(layer.builderHandle().arena.allocator(), op.result.dims, axis, batchDim(layer)),458            };459            break :blk .{ .value = try layer.next.bind(&batched_op, &.{}), .axis = axis };460        },461    };462}463464fn batchSelect(465    layer: anytype,466    op: *const program_mod.Operation,467    pred: BatchValue,468    on_true: BatchValue,469    on_false: BatchValue,470) !BatchValue {471    const batch_axis = pred.axis orelse on_true.axis orelse on_false.axis orelse {472        return .{473            .value = try layer.next.bind(op, &.{ pred.value, on_true.value, on_false.value }),474            .axis = null,475        };476    };477    const aligned_pred = try alignBatchOperand(layer, pred, batch_axis);478    const aligned_true = try alignBatchOperand(layer, on_true, batch_axis);479    const aligned_false = try alignBatchOperand(layer, on_false, batch_axis);480    return .{481        .value = try layer.next.bind(op, &.{ aligned_pred, aligned_true, aligned_false }),482        .axis = batch_axis,483    };484}485486fn alignBatchOperand(layer: anytype, operand: BatchValue, batch_axis: usize) !trace.Value {487    if (operand.axis) |axis| {488        if (axis == batch_axis) return operand.value;489        return moveBatchAxis(layer, operand.value, axis, batch_axis);490    }491    return broadcastToBatchAxis(layer, operand.value, batch_axis);492}493494fn broadcastToBatchAxis(layer: anytype, value: trace.Value, batch_axis: usize) !trace.Value {495    const allocator = layer.builderHandle().arena.allocator();496    const result_dims = try types.insertDim(allocator, value.ty.dims, batch_axis, batchDim(layer));497    const dims = try allocator.alloc(i64, value.ty.dims.len);498    var dim_index: usize = 0;499    for (0..result_dims.len) |axis| {500        if (axis == batch_axis) continue;501        dims[dim_index] = @intCast(axis);502        dim_index += 1;503    }504    var batched_op = program_mod.Operation{505        .id = program_mod.synthetic_id,506        .result = .{ .dtype = value.ty.dtype, .dims = result_dims },507        .kind = .{ .broadcast_in_dim = .{ .input = program_mod.synthetic_id, .broadcast_dims = dims } },508    };509    return layer.next.bind(&batched_op, &.{value});510}511512fn batchBinary(layer: anytype, op: *const program_mod.Operation, lhs: BatchValue, rhs: BatchValue) !BatchValue {513    if (lhs.axis == null and rhs.axis == null) {514        return .{ .value = try layer.next.bind(op, &.{ lhs.value, rhs.value }), .axis = null };515    }516517    if (lhs.axis) |lhs_axis| {518        if (rhs.axis) |rhs_axis| {519            const aligned_rhs = if (rhs_axis == lhs_axis) rhs.value else try moveBatchAxis(layer, rhs.value, rhs_axis, lhs_axis);520            return .{ .value = try layer.next.bind(op, &.{ lhs.value, aligned_rhs }), .axis = lhs_axis };521        }522        const lifted_rhs = try broadcastToBatchAxis(layer, rhs.value, lhs_axis);523        return .{ .value = try layer.next.bind(op, &.{ lhs.value, lifted_rhs }), .axis = lhs_axis };524    }525526    const rhs_axis = rhs.axis.?;527    const lifted_lhs = try broadcastToBatchAxis(layer, lhs.value, rhs_axis);528    return .{ .value = try layer.next.bind(op, &.{ lifted_lhs, rhs.value }), .axis = rhs_axis };529}530531fn batchBroadcast(532    layer: anytype,533    op: *const program_mod.Operation,534    broadcast: program_mod.Broadcast,535    input: BatchValue,536) !BatchValue {537    if (input.axis == null) return .{ .value = try layer.next.bind(op, &.{input.value}), .axis = null };538539    const moved = try moveBatchAxis(layer, input.value, input.axis.?, 0);540    if (input.value.ty.rank() != 1) return error.UnsupportedBatching;541    const result_dims = try types.insertDim(layer.builderHandle().arena.allocator(), op.result.dims, 0, batchDim(layer));542    var batched_op = op.*;543    batched_op.result = .{ .dtype = op.result.dtype, .dims = result_dims };544    batched_op.kind = .{ .broadcast_in_dim = .{ .input = broadcast.input, .broadcast_dims = &.{0} } };545    return .{ .value = try layer.next.bind(&batched_op, &.{moved}), .axis = 0 };546}547548fn batchBroadcastInDim(549    layer: anytype,550    op: *const program_mod.Operation,551    broadcast: program_mod.BroadcastInDim,552    input: BatchValue,553) !BatchValue {554    if (input.axis == null) {555        return .{ .value = try layer.next.bind(op, &.{input.value}), .axis = null };556    }557558    const moved = try moveBatchAxis(layer, input.value, input.axis.?, 0);559    const result_dims = try types.insertDim(layer.builderHandle().arena.allocator(), op.result.dims, 0, batchDim(layer));560    const dims = try shiftAxesWithLeading(layer.builderHandle().arena.allocator(), broadcast.broadcast_dims);561    var batched_op = op.*;562    batched_op.result = .{ .dtype = op.result.dtype, .dims = result_dims };563    batched_op.kind = .{ .broadcast_in_dim = .{ .input = broadcast.input, .broadcast_dims = dims } };564    return .{ .value = try layer.next.bind(&batched_op, &.{moved}), .axis = 0 };565}566567fn batchTranspose(568    layer: anytype,569    op: *const program_mod.Operation,570    transpose: program_mod.Transpose,571    input: BatchValue,572) !BatchValue {573    if (input.axis == null) return .{ .value = try layer.next.bind(op, &.{input.value}), .axis = null };574575    const allocator = layer.builderHandle().arena.allocator();576    const moved = try moveBatchAxis(layer, input.value, input.axis.?, 0);577    const permutation = try shiftAxesWithLeading(allocator, transpose.permutation);578    var batched_op = op.*;579    batched_op.result = .{ .dtype = input.value.ty.dtype, .dims = try types.permuted(allocator, moved.ty.dims, permutation) };580    batched_op.kind = .{ .transpose = .{ .input = transpose.input, .permutation = permutation } };581    return .{ .value = try layer.next.bind(&batched_op, &.{moved}), .axis = 0 };582}583584fn batchReduce(585    layer: anytype,586    op: *const program_mod.Operation,587    reduce: program_mod.Reduce,588    input: BatchValue,589    init: BatchValue,590) !BatchValue {591    if (input.axis == null and init.axis == null) {592        return .{ .value = try layer.next.bind(op, &.{ input.value, init.value }), .axis = null };593    }594    if (init.axis != null) return error.UnsupportedBatching;595    if (input.axis == null) return error.UnsupportedBatching;596597    const allocator = layer.builderHandle().arena.allocator();598    const moved = try moveBatchAxis(layer, input.value, input.axis.?, 0);599    const dimensions = try shiftAxes(allocator, reduce.dimensions, 1);600    var batched_op = op.*;601    batched_op.result = .{602        .dtype = op.result.dtype,603        .dims = try types.removeAxes(allocator, moved.ty.dims, dimensions),604    };605    batched_op.kind = .{ .reduce = .{ .input = reduce.input, .init = reduce.init, .reducer = reduce.reducer, .dimensions = dimensions } };606    return .{ .value = try layer.next.bind(&batched_op, &.{ moved, init.value }), .axis = 0 };607}608609fn batchDot(610    layer: anytype,611    op: *const program_mod.Operation,612    dot: program_mod.DotGeneral,613    lhs: BatchValue,614    rhs: BatchValue,615) !BatchValue {616    if (lhs.axis == null and rhs.axis == null) {617        return .{ .value = try layer.next.bind(op, &.{ lhs.value, rhs.value }), .axis = null };618    }619620    const allocator = layer.builderHandle().arena.allocator();621    const lhs_value = if (lhs.axis) |axis|622        try moveBatchAxis(layer, lhs.value, axis, 0)623    else624        try broadcastToBatchAxis(layer, lhs.value, 0);625    const rhs_value = if (rhs.axis) |axis|626        try moveBatchAxis(layer, rhs.value, axis, 0)627    else628        try broadcastToBatchAxis(layer, rhs.value, 0);629    const lhs_contract = try shiftAxes(allocator, dot.lhs_contract, 1);630    const rhs_contract = try shiftAxes(allocator, dot.rhs_contract, 1);631    const lhs_batch = try shiftAxesWithLeading(allocator, dot.lhs_batch);632    const rhs_batch = try shiftAxesWithLeading(allocator, dot.rhs_batch);633    const result_dims = try types.dotGeneralDims(634        allocator,635        lhs_value.ty.dims,636        rhs_value.ty.dims,637        lhs_contract,638        rhs_contract,639        lhs_batch,640        rhs_batch,641    );642    var batched_op = op.*;643    batched_op.result = .{ .dtype = op.result.dtype, .dims = result_dims };644    batched_op.kind = .{645        .dot_general = .{646            .lhs = dot.lhs,647            .rhs = dot.rhs,648            .lhs_contract = lhs_contract,649            .rhs_contract = rhs_contract,650            .lhs_batch = lhs_batch,651            .rhs_batch = rhs_batch,652        },653    };654    return .{ .value = try layer.next.bind(&batched_op, &.{ lhs_value, rhs_value }), .axis = 0 };655}656657fn emitUnbatched(layer: anytype, op: *const program_mod.Operation, map: []const BatchValue) !trace.Value {658    return switch (op.kind) {659        .parameter, .constant, .iota => layer.next.bind(op, &.{}),660        .unary => |unary| layer.next.bind(op, &.{map[unary.input.index].value}),661        .binary => |binary| layer.next.bind(op, &.{ map[binary.lhs.index].value, map[binary.rhs.index].value }),662        .broadcast => |broadcast| layer.next.bind(op, &.{map[broadcast.input.index].value}),663        .broadcast_in_dim => |broadcast| layer.next.bind(op, &.{map[broadcast.input.index].value}),664        .reshape => |reshape| layer.next.bind(op, &.{map[reshape.input.index].value}),665        .transpose => |transpose| layer.next.bind(op, &.{map[transpose.input.index].value}),666        .reduce => |reduce| layer.next.bind(op, &.{ map[reduce.input.index].value, map[reduce.init.index].value }),667        .gather => |gather| layer.next.bind(op, &.{ map[gather.input.index].value, map[gather.indices.index].value }),668        .scatter_add => |scatter_add| layer.next.bind(op, &.{669            map[scatter_add.input.index].value,670            map[scatter_add.indices.index].value,671            map[scatter_add.updates.index].value,672        }),673        .sparse_cross_entropy => |sparse_cross_entropy| layer.next.bind(op, &.{674            map[sparse_cross_entropy.logits.index].value,675            map[sparse_cross_entropy.targets.index].value,676        }),677        .compare => |compare| layer.next.bind(op, &.{ map[compare.lhs.index].value, map[compare.rhs.index].value }),678        .select => |select| layer.next.bind(op, &.{ map[select.pred.index].value, map[select.on_true.index].value, map[select.on_false.index].value }),679        .custom_call => error.CustomCallRequiresBatchContract,680        .dot_general => |dot| layer.next.bind(op, &.{ map[dot.lhs.index].value, map[dot.rhs.index].value }),681        .scan, .projection => error.ScanBatchingUnsupported,682    };683}684685fn broadcastOutputAxis(layer: anytype, value: trace.Value, options: Options) !trace.Value {686    return broadcastToBatchAxis(layer, value, options.out_axis);687}688689fn moveBatchAxis(layer: anytype, value: trace.Value, from: usize, to: usize) !trace.Value {690    if (from == to) return value;691    const rank = value.ty.rank();692    if (from >= rank or to >= rank) return error.AxisOutOfRange;693    const permutation = try layer.builderHandle().arena.allocator().alloc(i64, rank);694    var remaining: usize = 0;695    for (permutation, 0..) |*axis, index| {696        if (index == to) {697            axis.* = @intCast(from);698        } else {699            while (remaining == from) remaining += 1;700            axis.* = @intCast(remaining);701            remaining += 1;702        }703    }704    return emit.transpose(layer, value, permutation);705}706707fn shiftAxesWithLeading(allocator: std.mem.Allocator, axes: []const i64) ![]const i64 {708    const result = try allocator.alloc(i64, axes.len + 1);709    result[0] = 0;710    for (axes, 0..) |axis, index| {711        result[index + 1] = axis + 1;712    }713    return result;714}715716fn shiftAxes(allocator: std.mem.Allocator, axes: []const i64, amount: i64) ![]const i64 {717    const result = try allocator.alloc(i64, axes.len);718    for (axes, 0..) |axis, index| {719        result[index] = axis + amount;720    }721    return result;722}723724pub fn validate(source: *const program_mod.Program, options: Options) !void {725    if (options.axis_size < 0) return error.InvalidDimension;726    if (options.in_axes.len != source.parameters.len) return error.InAxisCountMismatch;727    for (source.parameters) |id| {728        const parameter = source.operation(id).kind.parameter;729        switch (options.in_axes[parameter.index]) {730            .none => {},731            .axis => |axis| if (axis > source.typeOf(id).rank()) return error.AxisOutOfRange,732        }733    }734}735736fn elementwiseBody(_: *trace.Builder, args: []const trace.Value) !trace.Value {737    return try (try args[0].mul(args[1])).tanh();738}739740fn expectExtents(expected: []const i64, ty: trace.Type) !void {741    try std.testing.expectEqual(expected.len, ty.dims.len);742    for (expected, ty.dims) |extent, dim| {743        try std.testing.expectEqual(extent, dim.extent);744    }745}746747test "tensor vmap batches elementwise programs with a generated axis name" {748    var source = try trace.define(std.testing.allocator, "vmap_elementwise", &.{749        types.spec(.f32, .{ .lane = 4 }),750        types.spec(.f32, .{ .lane = 4 }),751    }, elementwiseBody);752    defer source.deinit();753754    var batched = try vmap(std.testing.allocator, &source, .{755        .axis_size = 8,756        .in_axes = &.{ mapped(0), mapped(0) },757    });758    defer batched.deinit();759760    try std.testing.expectEqual(@as(usize, 2), batched.parameters.len);761    try expectExtents(&.{ 8, 4 }, batched.typeOf(batched.parameters[0]));762    try expectExtents(&.{ 8, 4 }, batched.typeOf(batched.outputs[0]));763    try std.testing.expectEqualStrings("#batch", batched.typeOf(batched.outputs[0]).dims[0].name);764    try std.testing.expectEqualStrings("lane", batched.typeOf(batched.outputs[0]).dims[1].name);765}766767test "tensor vmap names the batch axis on request" {768    var source = try trace.define(std.testing.allocator, "vmap_named_axis", &.{769        types.spec(.f32, .{ .lane = 4 }),770        types.spec(.f32, .{ .lane = 4 }),771    }, elementwiseBody);772    defer source.deinit();773774    var batched = try vmap(std.testing.allocator, &source, .{775        .axis_size = 8,776        .in_axes = &.{ mapped(0), mapped(0) },777        .axis_name = "walk",778    });779    defer batched.deinit();780781    try std.testing.expectEqualStrings("walk", batched.typeOf(batched.outputs[0]).dims[0].name);782}783784test "tensor vmap generates fresh names for nested batching" {785    var source = try trace.define(std.testing.allocator, "vmap_nested", &.{786        types.spec(.f32, .{ .lane = 4 }),787        types.spec(.f32, .{ .lane = 4 }),788    }, elementwiseBody);789    defer source.deinit();790791    var once = try vmap(std.testing.allocator, &source, .{792        .axis_size = 8,793        .in_axes = &.{ mapped(0), mapped(0) },794    });795    defer once.deinit();796797    var twice = try vmap(std.testing.allocator, &once, .{798        .axis_size = 3,799        .in_axes = &.{ mapped(0), mapped(0) },800    });801    defer twice.deinit();802803    const out = twice.typeOf(twice.outputs[0]);804    try expectExtents(&.{ 3, 8, 4 }, out);805    try std.testing.expectEqualStrings("#batch1", out.dims[0].name);806    try std.testing.expectEqualStrings("#batch", out.dims[1].name);807}808809fn unmappedOutputBody(_: *trace.Builder, args: []const trace.Value) !trace.Value {810    return args[1];811}812813test "tensor vmap broadcasts unmapped outputs" {814    var source = try trace.define(std.testing.allocator, "vmap_unmapped_output", &.{815        types.spec(.f32, .{ .lane = 4 }),816        types.spec(.f32, .{ .lane = 4 }),817    }, unmappedOutputBody);818    defer source.deinit();819820    var batched = try vmap(std.testing.allocator, &source, .{821        .axis_size = 8,822        .in_axes = &.{ mapped(0), .none },823    });824    defer batched.deinit();825826    try std.testing.expectEqual(@as(usize, 2), batched.parameters.len);827    try expectExtents(&.{ 8, 4 }, batched.typeOf(batched.outputs[0]));828}829830const GeneratedCounts = struct {831    broadcast_in_dim: usize = 0,832    transpose: usize = 0,833};834835const GeneratedCounter = struct {836    counts: *GeneratedCounts,837838    pub fn broadcastInDim(self: *@This(), ctx: anytype) !trace.Value {839        self.counts.broadcast_in_dim += 1;840        return ctx.default();841    }842843    pub fn transpose(self: *@This(), ctx: anytype) !trace.Value {844        self.counts.transpose += 1;845        return ctx.default();846    }847};848849test "tensor vmap binds generated output broadcasts through downstream semantics" {850    var source = try trace.define(std.testing.allocator, "vmap_generated_broadcast", &.{851        types.spec(.f32, .{ .lane = 4 }),852        types.spec(.f32, .{ .lane = 4 }),853    }, unmappedOutputBody);854    defer source.deinit();855856    var builder = try trace.Builder.init(std.testing.allocator, source.name);857    errdefer builder.deinit();858859    var counts: GeneratedCounts = .{};860    const graph = interpret.Graph{ .builder = &builder };861    const counted = interpret.bind(GeneratedCounter{ .counts = &counts }).attach(graph);862    var batched = try interpret.run(std.testing.allocator, &source, semantics(counted, .{863        .axis_size = 8,864        .in_axes = &.{ mapped(0), .none },865    }));866    defer batched.deinit();867868    try std.testing.expectEqual(@as(usize, 1), counts.broadcast_in_dim);869    try std.testing.expectEqual(@as(usize, 0), counts.transpose);870    try expectExtents(&.{ 8, 4 }, batched.typeOf(batched.outputs[0]));871}872873test "tensor vmapWith routes batch generated ops through user semantics" {874    var source = try trace.define(std.testing.allocator, "vmap_with_generated_broadcast", &.{875        types.spec(.f32, .{ .lane = 4 }),876        types.spec(.f32, .{ .lane = 4 }),877    }, unmappedOutputBody);878    defer source.deinit();879880    var counts: GeneratedCounts = .{};881    var batched = try vmapWith(882        std.testing.allocator,883        &source,884        .{885            .axis_size = 8,886            .in_axes = &.{ mapped(0), .none },887        },888        .{ .batch = interpret.bind(GeneratedCounter{ .counts = &counts }) },889    );890    defer batched.deinit();891892    try std.testing.expectEqual(@as(usize, 1), counts.broadcast_in_dim);893    try std.testing.expectEqual(@as(usize, 0), counts.transpose);894    try expectExtents(&.{ 8, 4 }, batched.typeOf(batched.outputs[0]));895}896897const GeneratedBroadcastRewrite = struct {898    seen: *usize,899900    pub fn broadcastInDim(self: *@This(), ctx: *transform.Context) !?trace.Value {901        self.seen.* += 1;902        try std.testing.expect(!ctx.isZero(0));903        try std.testing.expect(ctx.constantPayload(0) == null);904        return null;905    }906};907908test "tensor vmap generated ops are safe for downstream rewrite metadata queries" {909    var source = try trace.define(std.testing.allocator, "vmap_generated_rewrite_metadata", &.{910        types.spec(.f32, .{ .lane = 4 }),911        types.spec(.f32, .{ .lane = 4 }),912    }, unmappedOutputBody);913    defer source.deinit();914915    var builder = try trace.Builder.init(std.testing.allocator, source.name);916    errdefer builder.deinit();917918    var seen: usize = 0;919    const graph = interpret.Graph{ .builder = &builder };920    const rewrite = transform.semantics(&source, graph, GeneratedBroadcastRewrite{ .seen = &seen });921    var batched = try interpret.run(std.testing.allocator, &source, semantics(rewrite, .{922        .axis_size = 8,923        .in_axes = &.{ mapped(0), .none },924    }));925    defer batched.deinit();926927    try std.testing.expectEqual(@as(usize, 1), seen);928    try expectExtents(&.{ 8, 4 }, batched.typeOf(batched.outputs[0]));929}930931fn gatherBody(_: *trace.Builder, args: []const trace.Value) !trace.Value {932    return args[0].gather(args[1], .vocab);933}934935test "tensor vmap gathers with batched indices" {936    var source = try trace.define(std.testing.allocator, "vmap_gather_indices", &.{937        types.spec(.f32, .{ .vocab = 32, .channel = 8 }),938        types.spec(.i32, .{ .token = 5 }),939    }, gatherBody);940    defer source.deinit();941942    var batched = try vmap(std.testing.allocator, &source, .{943        .axis_size = 7,944        .in_axes = &.{ .none, mapped(0) },945    });946    defer batched.deinit();947948    try expectExtents(&.{ 32, 8 }, batched.typeOf(batched.parameters[0]));949    try expectExtents(&.{ 7, 5 }, batched.typeOf(batched.parameters[1]));950    try expectExtents(&.{ 7, 5, 8 }, batched.typeOf(batched.outputs[0]));951    try std.testing.expectEqualStrings("#batch", batched.typeOf(batched.outputs[0]).dims[0].name);952    try std.testing.expectEqualStrings("token", batched.typeOf(batched.outputs[0]).dims[1].name);953}954955test "tensor vmap gathers with a batched input" {956    var source = try trace.define(std.testing.allocator, "vmap_gather_input", &.{957        types.spec(.f32, .{ .vocab = 32, .channel = 8 }),958        types.spec(.i32, .{ .token = 5 }),959    }, gatherBody);960    defer source.deinit();961962    var batched = try vmap(std.testing.allocator, &source, .{963        .axis_size = 7,964        .in_axes = &.{ mapped(0), .none },965    });966    defer batched.deinit();967968    try expectExtents(&.{ 7, 32, 8 }, batched.typeOf(batched.parameters[0]));969    try expectExtents(&.{5}, batched.typeOf(batched.parameters[1]));970    try expectExtents(&.{ 7, 5, 8 }, batched.typeOf(batched.outputs[0]));971    try std.testing.expectEqualStrings("#batch", batched.typeOf(batched.outputs[0]).dims[0].name);972    try std.testing.expectEqualStrings("token", batched.typeOf(batched.outputs[0]).dims[1].name);973}974975test "tensor vmap gathers with shared batched input and indices" {976    var source = try trace.define(std.testing.allocator, "vmap_gather_shared_batch", &.{977        types.spec(.f32, .{ .vocab = 32, .channel = 8 }),978        types.spec(.i32, .{ .token = 5 }),979    }, gatherBody);980    defer source.deinit();981982    var batched = try vmap(std.testing.allocator, &source, .{983        .axis_size = 7,984        .in_axes = &.{ mapped(0), mapped(0) },985    });986    defer batched.deinit();987988    try expectExtents(&.{ 7, 32, 8 }, batched.typeOf(batched.parameters[0]));989    try expectExtents(&.{ 7, 5 }, batched.typeOf(batched.parameters[1]));990    try expectExtents(&.{ 7, 5, 8 }, batched.typeOf(batched.outputs[0]));991    try std.testing.expectEqualStrings("#batch", batched.typeOf(batched.outputs[0]).dims[0].name);992    try std.testing.expectEqualStrings("token", batched.typeOf(batched.outputs[0]).dims[1].name);993}994995fn scatterAddBody(_: *trace.Builder, args: []const trace.Value) !trace.Value {996    return args[0].scatterAdd(args[1], args[2], .vocab);997}998999test "tensor vmap scatters with batched updates and shared indices" {1000    var source = try trace.define(std.testing.allocator, "vmap_scatter_add_updates", &.{1001        types.spec(.f32, .{ .vocab = 32, .channel = 8 }),1002        types.spec(.i32, .{ .token = 5 }),1003        types.spec(.f32, .{ .token = 5, .channel = 8 }),1004    }, scatterAddBody);1005    defer source.deinit();10061007    var batched = try vmap(std.testing.allocator, &source, .{1008        .axis_size = 7,1009        .in_axes = &.{ mapped(0), .none, mapped(0) },1010    });1011    defer batched.deinit();10121013    try expectExtents(&.{ 7, 32, 8 }, batched.typeOf(batched.parameters[0]));1014    try expectExtents(&.{5}, batched.typeOf(batched.parameters[1]));1015    try expectExtents(&.{ 7, 5, 8 }, batched.typeOf(batched.parameters[2]));1016    try expectExtents(&.{ 7, 32, 8 }, batched.typeOf(batched.outputs[0]));1017    try std.testing.expectEqualStrings("#batch", batched.typeOf(batched.outputs[0]).dims[0].name);1018}10191020test "tensor vmap scatters with shared batched indices" {1021    var source = try trace.define(std.testing.allocator, "vmap_scatter_add_shared_indices", &.{1022        types.spec(.f32, .{ .vocab = 32, .channel = 8 }),1023        types.spec(.i32, .{ .token = 5 }),1024        types.spec(.f32, .{ .token = 5, .channel = 8 }),1025    }, scatterAddBody);1026    defer source.deinit();10271028    var batched = try vmap(std.testing.allocator, &source, .{1029        .axis_size = 7,1030        .in_axes = &.{ mapped(0), mapped(0), mapped(0) },1031    });1032    defer batched.deinit();10331034    try expectExtents(&.{ 7, 32, 8 }, batched.typeOf(batched.parameters[0]));1035    try expectExtents(&.{ 7, 5 }, batched.typeOf(batched.parameters[1]));1036    try expectExtents(&.{ 7, 5, 8 }, batched.typeOf(batched.parameters[2]));1037    try expectExtents(&.{ 7, 32, 8 }, batched.typeOf(batched.outputs[0]));1038    try std.testing.expectEqualStrings("#batch", batched.typeOf(batched.outputs[0]).dims[0].name);1039}10401041test "tensor vmap scatters with batched indices and shared input" {1042    var source = try trace.define(std.testing.allocator, "vmap_scatter_add_batched_indices_shared_input", &.{1043        types.spec(.f32, .{ .vocab = 32, .channel = 8 }),1044        types.spec(.i32, .{ .token = 5 }),1045        types.spec(.f32, .{ .token = 5, .channel = 8 }),1046    }, scatterAddBody);1047    defer source.deinit();10481049    var batched = try vmap(std.testing.allocator, &source, .{1050        .axis_size = 7,1051        .in_axes = &.{ .none, mapped(0), mapped(0) },1052    });1053    defer batched.deinit();10541055    try expectExtents(&.{ 32, 8 }, batched.typeOf(batched.parameters[0]));1056    try expectExtents(&.{ 7, 5 }, batched.typeOf(batched.parameters[1]));1057    try expectExtents(&.{ 7, 5, 8 }, batched.typeOf(batched.parameters[2]));1058    try expectExtents(&.{ 7, 32, 8 }, batched.typeOf(batched.outputs[0]));1059    try std.testing.expectEqualStrings("#batch", batched.typeOf(batched.outputs[0]).dims[0].name);1060}10611062fn identityBody(_: *trace.Builder, args: []const trace.Value) !trace.Value {1063    return args[0];1064}10651066test "tensor vmap binds generated batch-axis moves through downstream semantics" {1067    var source = try trace.define(std.testing.allocator, "vmap_generated_transpose", &.{1068        types.spec(.f32, .{ .row = 2, .col = 4 }),1069    }, identityBody);1070    defer source.deinit();10711072    var builder = try trace.Builder.init(std.testing.allocator, source.name);1073    errdefer builder.deinit();10741075    var counts: GeneratedCounts = .{};1076    const graph = interpret.Graph{ .builder = &builder };1077    const counted = interpret.bind(GeneratedCounter{ .counts = &counts }).attach(graph);1078    var batched = try interpret.run(std.testing.allocator, &source, semantics(counted, .{1079        .axis_size = 8,1080        .in_axes = &.{mapped(1)},1081        .out_axis = 0,1082    }));1083    defer batched.deinit();10841085    try std.testing.expectEqual(@as(usize, 0), counts.broadcast_in_dim);1086    try std.testing.expectEqual(@as(usize, 1), counts.transpose);1087    try expectExtents(&.{ 8, 2, 4 }, batched.typeOf(batched.outputs[0]));1088}10891090fn contractBody(_: *trace.Builder, args: []const trace.Value) !trace.Value {1091    return try args[0].contract(args[1], .k);1092}10931094test "tensor vmap batches contraction with an unmapped rhs" {1095    var source = try trace.define(std.testing.allocator, "vmap_contract", &.{1096        types.spec(.f32, .{ .m = 2, .k = 4 }),1097        types.spec(.f32, .{ .k = 4, .n = 3 }),1098    }, contractBody);1099    defer source.deinit();11001101    var batched = try vmap(std.testing.allocator, &source, .{1102        .axis_size = 8,1103        .in_axes = &.{ mapped(0), .none },1104    });1105    defer batched.deinit();11061107    try std.testing.expectEqual(@as(usize, 2), batched.parameters.len);1108    try expectExtents(&.{ 8, 2, 3 }, batched.typeOf(batched.outputs[0]));1109    try std.testing.expectEqualStrings("#batch", batched.typeOf(batched.outputs[0]).dims[0].name);1110}11111112fn addZeroBody(builder: *trace.Builder, args: []const trace.Value) !trace.Value {1113    const zero = try builder.full(.f32, .{ .lane = 4 }, 0.0);1114    return try args[0].add(zero);1115}11161117const BatchDropAddZero = struct {1118    pub fn add(_: *@This(), ctx: *transform.Context) !?trace.Value {1119        if (ctx.isZero(1)) return ctx.arg(0);1120        if (ctx.isZero(0)) return ctx.arg(1);1121        return null;1122    }1123};11241125test "tensor vmap can delegate primitive emission through rewrite semantics" {1126    var source = try trace.define(std.testing.allocator, "vmap_rewrite", &.{1127        types.spec(.f32, .{ .lane = 4 }),1128    }, addZeroBody);1129    defer source.deinit();11301131    const options = Options{1132        .axis_size = 8,1133        .in_axes = &.{mapped(0)},1134    };11351136    var plain = try vmap(std.testing.allocator, &source, options);1137    defer plain.deinit();11381139    var builder = try trace.Builder.init(std.testing.allocator, source.name);1140    errdefer builder.deinit();1141    const graph = interpret.Graph{ .builder = &builder };1142    const rewrite = transform.semantics(&source, graph, BatchDropAddZero{});1143    var rewritten = try interpret.run(std.testing.allocator, &source, semantics(rewrite, options));1144    defer rewritten.deinit();11451146    try std.testing.expect(rewritten.operationCount() < plain.operationCount());1147    try std.testing.expectEqual(@as(u32, 0), rewritten.outputs[0].index);1148    try expectExtents(&.{ 8, 4 }, rewritten.typeOf(rewritten.outputs[0]));1149}

1pub const batch = @import("batch.zig");