lib/accy/src/tensor/batch.zig

daab053ee43316e1809a84551d573ddd1e5bf3d2

   1 const std = @import("std");
   2 const emit = @import("emit.zig");
   3 const hook = @import("hook.zig");
   4 const interpret = @import("interpret/root.zig");
   5 const program_mod = @import("program.zig");
   6 const trace = @import("trace/root.zig");
   7 const transform = @import("transform.zig");
   8 const types = @import("type/root.zig");
   9 
  10 pub const Axis = union(enum) {
  11     none,
  12     axis: usize,
  13 };
  14 
  15 pub const Options = struct {
  16     axis_size: i64,
  17     in_axes: []const Axis,
  18     out_axis: usize = 0,
  19     axis_name: []const u8 = &.{},
  20 };
  21 
  22 const default_axis_name = "#batch";
  23 
  24 const BatchValue = struct {
  25     value: trace.Value,
  26     axis: ?usize,
  27 };
  28 
  29 pub fn mapped(axis_index: usize) Axis {
  30     return .{ .axis = axis_index };
  31 }
  32 
  33 pub 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);
  39 
  40     var builder = try trace.Builder.init(allocator, source.name);
  41     errdefer builder.deinit();
  42 
  43     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 }
  47 
  48 pub 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);
  55 
  56     var builder = try trace.Builder.init(allocator, source.name);
  57     errdefer builder.deinit();
  58 
  59     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 }
  63 
  64 fn 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 }
  79 
  80 fn 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 }
  89 
  90 fn 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 }
  99 
 100 pub fn semantics(next: anytype, options: Options) Semantics(@TypeOf(next)) {
 101     return .{
 102         .next = next,
 103         .options = options,
 104     };
 105 }
 106 
 107 pub 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 }
 110 
 111 pub fn Semantics(comptime Next: type) type {
 112     return struct {
 113         next: Next,
 114         options: Options,
 115 
 116         pub const Value: type = BatchValue;
 117         pub const Result: type = Next.Result;
 118 
 119         pub fn operation(self: *@This(), step: *interpret.Step(Value)) !Value {
 120             return batchOperation(self, step.op, step.values);
 121         }
 122 
 123         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                 else
 129                     try broadcastOutputAxis(self, output.value, self.options);
 130             }
 131             return self.next.finish(values);
 132         }
 133 
 134         pub fn builderHandle(self: *@This()) *trace.Builder {
 135             return self.next.builderHandle();
 136         }
 137 
 138         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 }
 144 
 145 fn batchDim(layer: anytype) types.Dim {
 146     return .{ .name = layer.axisName(), .extent = layer.options.axis_size };
 147 }
 148 
 149 fn 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 }
 199 
 200 fn 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     }
 210 
 211     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         else
 218             try broadcastToBatchAxis(layer, carried.value, 0);
 219     }
 220 
 221     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();
 230 
 231     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();
 239 
 240     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 }
 253 
 254 fn 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 }
 280 
 281 fn 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     else
 289         try broadcastToBatchAxis(layer, input.value, 0);
 290     const moved_updates = if (updates.axis) |axis|
 291         try moveBatchAxis(layer, updates.value, axis, 0)
 292     else
 293         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 }
 306 
 307 fn 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     else
 315         try broadcastToBatchAxis(layer, logits.value, 0);
 316     const moved_targets = if (targets.axis) |axis|
 317         try moveBatchAxis(layer, targets.value, axis, 0)
 318     else
 319         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 }
 325 
 326 fn 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 }
 341 
 342 fn 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);
 354 
 355     const reduced = if (index_dims.len == 0)
 356         selected
 357     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 }
 365 
 366 fn 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 }
 374 
 375 fn 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 }
 398 
 399 fn 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         else
 406             @intCast(index + index_rank);
 407     }
 408     return emit.broadcastInDim(layer, input, expanded_dims, mapping);
 409 }
 410 
 411 fn 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 }
 418 
 419 fn 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         else
 426             @intCast(index + 1);
 427     }
 428     return emit.broadcastInDim(layer, updates, expanded_dims, mapping);
 429 }
 430 
 431 fn 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 }
 435 
 436 fn 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 }
 444 
 445 fn leadingAxis(axis: i64) !i64 {
 446     if (axis < 0) return error.AxisOutOfRange;
 447     return std.math.add(i64, axis, 1) catch error.AxisOutOfRange;
 448 }
 449 
 450 fn 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 }
 463 
 464 fn 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 }
 485 
 486 fn 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 }
 493 
 494 fn 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 }
 511 
 512 fn 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     }
 516 
 517     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     }
 525 
 526     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 }
 530 
 531 fn 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 };
 538 
 539     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 }
 547 
 548 fn 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     }
 557 
 558     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 }
 566 
 567 fn 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 };
 574 
 575     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 }
 583 
 584 fn 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;
 596 
 597     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 }
 608 
 609 fn 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     }
 619 
 620     const allocator = layer.builderHandle().arena.allocator();
 621     const lhs_value = if (lhs.axis) |axis|
 622         try moveBatchAxis(layer, lhs.value, axis, 0)
 623     else
 624         try broadcastToBatchAxis(layer, lhs.value, 0);
 625     const rhs_value = if (rhs.axis) |axis|
 626         try moveBatchAxis(layer, rhs.value, axis, 0)
 627     else
 628         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 }
 656 
 657 fn 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 }
 684 
 685 fn broadcastOutputAxis(layer: anytype, value: trace.Value, options: Options) !trace.Value {
 686     return broadcastToBatchAxis(layer, value, options.out_axis);
 687 }
 688 
 689 fn 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 }
 706 
 707 fn 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 }
 715 
 716 fn 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 }
 723 
 724 pub 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 }
 735 
 736 fn elementwiseBody(_: *trace.Builder, args: []const trace.Value) !trace.Value {
 737     return try (try args[0].mul(args[1])).tanh();
 738 }
 739 
 740 fn 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 }
 746 
 747 test "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();
 753 
 754     var batched = try vmap(std.testing.allocator, &source, .{
 755         .axis_size = 8,
 756         .in_axes = &.{ mapped(0), mapped(0) },
 757     });
 758     defer batched.deinit();
 759 
 760     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 }
 766 
 767 test "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();
 773 
 774     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();
 780 
 781     try std.testing.expectEqualStrings("walk", batched.typeOf(batched.outputs[0]).dims[0].name);
 782 }
 783 
 784 test "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();
 790 
 791     var once = try vmap(std.testing.allocator, &source, .{
 792         .axis_size = 8,
 793         .in_axes = &.{ mapped(0), mapped(0) },
 794     });
 795     defer once.deinit();
 796 
 797     var twice = try vmap(std.testing.allocator, &once, .{
 798         .axis_size = 3,
 799         .in_axes = &.{ mapped(0), mapped(0) },
 800     });
 801     defer twice.deinit();
 802 
 803     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 }
 808 
 809 fn unmappedOutputBody(_: *trace.Builder, args: []const trace.Value) !trace.Value {
 810     return args[1];
 811 }
 812 
 813 test "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();
 819 
 820     var batched = try vmap(std.testing.allocator, &source, .{
 821         .axis_size = 8,
 822         .in_axes = &.{ mapped(0), .none },
 823     });
 824     defer batched.deinit();
 825 
 826     try std.testing.expectEqual(@as(usize, 2), batched.parameters.len);
 827     try expectExtents(&.{ 8, 4 }, batched.typeOf(batched.outputs[0]));
 828 }
 829 
 830 const GeneratedCounts = struct {
 831     broadcast_in_dim: usize = 0,
 832     transpose: usize = 0,
 833 };
 834 
 835 const GeneratedCounter = struct {
 836     counts: *GeneratedCounts,
 837 
 838     pub fn broadcastInDim(self: *@This(), ctx: anytype) !trace.Value {
 839         self.counts.broadcast_in_dim += 1;
 840         return ctx.default();
 841     }
 842 
 843     pub fn transpose(self: *@This(), ctx: anytype) !trace.Value {
 844         self.counts.transpose += 1;
 845         return ctx.default();
 846     }
 847 };
 848 
 849 test "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();
 855 
 856     var builder = try trace.Builder.init(std.testing.allocator, source.name);
 857     errdefer builder.deinit();
 858 
 859     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();
 867 
 868     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 }
 872 
 873 test "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();
 879 
 880     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();
 891 
 892     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 }
 896 
 897 const GeneratedBroadcastRewrite = struct {
 898     seen: *usize,
 899 
 900     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 };
 907 
 908 test "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();
 914 
 915     var builder = try trace.Builder.init(std.testing.allocator, source.name);
 916     errdefer builder.deinit();
 917 
 918     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();
 926 
 927     try std.testing.expectEqual(@as(usize, 1), seen);
 928     try expectExtents(&.{ 8, 4 }, batched.typeOf(batched.outputs[0]));
 929 }
 930 
 931 fn gatherBody(_: *trace.Builder, args: []const trace.Value) !trace.Value {
 932     return args[0].gather(args[1], .vocab);
 933 }
 934 
 935 test "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();
 941 
 942     var batched = try vmap(std.testing.allocator, &source, .{
 943         .axis_size = 7,
 944         .in_axes = &.{ .none, mapped(0) },
 945     });
 946     defer batched.deinit();
 947 
 948     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 }
 954 
 955 test "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();
 961 
 962     var batched = try vmap(std.testing.allocator, &source, .{
 963         .axis_size = 7,
 964         .in_axes = &.{ mapped(0), .none },
 965     });
 966     defer batched.deinit();
 967 
 968     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 }
 974 
 975 test "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();
 981 
 982     var batched = try vmap(std.testing.allocator, &source, .{
 983         .axis_size = 7,
 984         .in_axes = &.{ mapped(0), mapped(0) },
 985     });
 986     defer batched.deinit();
 987 
 988     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 }
 994 
 995 fn scatterAddBody(_: *trace.Builder, args: []const trace.Value) !trace.Value {
 996     return args[0].scatterAdd(args[1], args[2], .vocab);
 997 }
 998 
 999 test "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();
1006 
1007     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();
1012 
1013     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 }
1019 
1020 test "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();
1027 
1028     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();
1033 
1034     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 }
1040 
1041 test "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();
1048 
1049     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();
1054 
1055     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 }
1061 
1062 fn identityBody(_: *trace.Builder, args: []const trace.Value) !trace.Value {
1063     return args[0];
1064 }
1065 
1066 test "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();
1071 
1072     var builder = try trace.Builder.init(std.testing.allocator, source.name);
1073     errdefer builder.deinit();
1074 
1075     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();
1084 
1085     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 }
1089 
1090 fn contractBody(_: *trace.Builder, args: []const trace.Value) !trace.Value {
1091     return try args[0].contract(args[1], .k);
1092 }
1093 
1094 test "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();
1100 
1101     var batched = try vmap(std.testing.allocator, &source, .{
1102         .axis_size = 8,
1103         .in_axes = &.{ mapped(0), .none },
1104     });
1105     defer batched.deinit();
1106 
1107     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 }
1111 
1112 fn 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 }
1116 
1117 const 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 };
1124 
1125 test "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();
1130 
1131     const options = Options{
1132         .axis_size = 8,
1133         .in_axes = &.{mapped(0)},
1134     };
1135 
1136     var plain = try vmap(std.testing.allocator, &source, options);
1137     defer plain.deinit();
1138 
1139     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();
1145 
1146     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 }