tiny.accy.tensor.trace.builder

1const std = @import("std");2const tensor = @import("../root.zig");3const trace = @import("root.zig");45const Binary = trace.Binary;6const CompareDirection = trace.CompareDirection;7const DType = trace.DType;8const Dim = trace.Dim;9const Id = trace.Id;10const Program = trace.Program;11const Reducer = trace.Reducer;12const Spec = trace.Spec;13const Type = trace.Type;14const Unary = trace.Unary;15const Value = trace.Value;16const program_mod = tensor.program;17const type_mod = tensor.types;1819pub const ScanScope = struct {20    parent: *Builder,21    child: Builder,22    length: i64,23    init_ids: []const Id,24    phase: Phase = .open,2526    const Phase = enum { open, done };2728    pub fn body(self: *ScanScope) *Builder {29        return &self.child;30    }3132    pub fn carry(self: *ScanScope, index: usize) Value {33        return self.child.makeValue(self.child.parameters_list.items[index]);34    }3536    pub fn finish(self: *ScanScope, next: []const Value) !Value {37        if (self.phase != .open) return error.ScanScopeMisused;38        if (next.len != self.init_ids.len) return error.ScanArityMismatch;39        const output_ids = try self.child.arena.allocator().alloc(Id, next.len);40        for (next, output_ids) |value, *slot| {41            try self.child.ensureValue(value);42            slot.* = value.id;43        }44        const body_view = program_mod.Subgraph{45            .values = self.child.values.items,46            .operations = self.child.operations.items,47            .parameters = self.child.parameters_list.items,48            .outputs = output_ids,49        };50        const result = try self.parent.emitScanIds(self.length, self.init_ids, &body_view);51        self.abort();52        return result;53    }5455    pub fn abort(self: *ScanScope) void {56        if (self.phase == .done) return;57        self.child.deinit();58        self.phase = .done;59    }60};6162pub const Builder = struct {63    allocator: std.mem.Allocator,64    arena: std.heap.ArenaAllocator,65    name: []const u8,66    values: std.ArrayListUnmanaged(Type) = .empty,67    operations: std.ArrayListUnmanaged(program_mod.Operation) = .empty,68    parameters_list: std.ArrayListUnmanaged(Id) = .empty,69    finished: bool = false,7071    pub fn init(allocator: std.mem.Allocator, name: []const u8) !Builder {72        var arena = std.heap.ArenaAllocator.init(allocator);73        errdefer arena.deinit();74        var builder = Builder{75            .allocator = allocator,76            .arena = arena,77            .name = &.{},78        };79        builder.name = try builder.arena.allocator().dupe(u8, name);80        return builder;81    }8283    pub fn deinit(self: *Builder) void {84        if (!self.finished) {85            self.arena.deinit();86            self.finished = true;87        }88    }8990    pub fn input(self: *Builder, dtype: DType, dims_struct: anytype) !Value {91        var buffer: [type_mod.dimCount(@TypeOf(dims_struct))]Dim = undefined;92        type_mod.fillDims(dims_struct, &buffer);93        return self.inputDims(dtype, &buffer);94    }9596    pub fn inputDims(self: *Builder, dtype: DType, dims: []const Dim) !Value {97        try type_mod.validateAuthoredDims(dims);98        return self.inputTyped(.{ .dtype = dtype, .dims = dims });99    }100101    pub fn inputSpec(self: *Builder, input_spec: Spec) !Value {102        return self.inputDims(input_spec.dtype, input_spec.dims);103    }104105    pub fn inputs(self: *Builder, specs: []const Spec) ![]Value {106        const result = try self.arena.allocator().alloc(Value, specs.len);107        for (specs, 0..) |input_spec, index| {108            result[index] = try self.inputSpec(input_spec);109        }110        return result;111    }112113    pub fn inputTyped(self: *Builder, ty: Type) !Value {114        const owned = try self.copyType(ty);115        const id = try self.append(116            owned,117            .{ .parameter = .{ .index = self.parameters_list.items.len } },118        );119        try self.parameters_list.append(self.arena.allocator(), id);120        return self.makeValue(id);121    }122123    pub fn operation(self: *Builder, op: *const program_mod.Operation, args: []const Value) !Value {124        return switch (op.kind) {125            .parameter => self.inputTyped(op.result),126            .constant => |constant| self.constantBytes(op.result, constant.payload),127            .iota => |iota_op| self.iotaTyped(op.result, iota_op.axis),128            .unary => |unary_op| self.unary(unary_op.op, args[0]),129            .binary => |binary_op| self.binary(binary_op.op, args[0], args[1]),130            .broadcast => self.broadcastOp(args[0], op.result.dims),131            .broadcast_in_dim => |broadcast_op| self.broadcastTo(args[0], op.result.dims, broadcast_op.broadcast_dims),132            .reshape => self.reshapeTo(args[0], op.result.dims),133            .transpose => |transpose_op| self.transposeBy(args[0], transpose_op.permutation),134            .reduce => |reduce_op| self.reduce(args[0], args[1], reduce_op.reducer, reduce_op.dimensions),135            .gather => |gather_op| self.gatherOp(args[0], args[1], gather_op.axis),136            .scatter_add => |scatter_add| self.scatterAddOp(args[0], args[1], args[2], scatter_add.axis),137            .sparse_cross_entropy => |sparse_cross_entropy_op| self.sparseCrossEntropyOp(args[0], args[1], sparse_cross_entropy_op.axis),138            .compare => |compare_op| self.compare(compare_op.direction, args[0], args[1]),139            .select => self.select(args[0], args[1], args[2]),140            .custom_call => |custom_call| self.customCall(custom_call.target, custom_call.version, args, op.result),141            .dot_general => |dot| self.dotGeneralOp(142                args[0],143                args[1],144                dot.lhs_contract,145                dot.rhs_contract,146                dot.lhs_batch,147                dot.rhs_batch,148            ),149            .scan => |scan_op| self.emitScan(scan_op.length, args[0..scan_op.inits.len], scan_op.body),150            .projection => |projection_op| self.projection(args[0], projection_op.index),151        };152    }153154    pub fn scalar(self: *Builder, comptime dtype: DType, scalar_value: dtype.ZigType()) !Value {155        const ty = Type.scalar(dtype);156        return self.constantBytes(ty, std.mem.asBytes(&scalar_value));157    }158159    pub fn full(self: *Builder, comptime dtype: DType, dims_struct: anytype, fill_value: dtype.ZigType()) !Value {160        var buffer: [type_mod.dimCount(@TypeOf(dims_struct))]Dim = undefined;161        type_mod.fillDims(dims_struct, &buffer);162        return self.fullDims(dtype, &buffer, fill_value);163    }164165    pub fn fullDims(self: *Builder, comptime dtype: DType, dims: []const Dim, fill_value: dtype.ZigType()) !Value {166        try type_mod.validateAuthoredDims(dims);167        const ty = try Type.init(self.arena.allocator(), dtype, dims);168        const count = try ty.elementCount();169        const payload = try self.arena.allocator().alloc(dtype.ZigType(), count);170        for (payload) |*slot| {171            slot.* = fill_value;172        }173        return self.constantOwned(ty, std.mem.sliceAsBytes(payload));174    }175176    pub fn fullFloat(self: *Builder, ty: Type, fill_value: f64) !Value {177        const owned = try self.copyType(ty);178        const count = try owned.elementCount();179        return switch (owned.dtype) {180            .f16 => blk: {181                const payload = try self.arena.allocator().alloc(f16, count);182                const value: f16 = @floatCast(fill_value);183                for (payload) |*slot| slot.* = value;184                break :blk self.constantOwned(owned, std.mem.sliceAsBytes(payload));185            },186            .bf16 => blk: {187                const Bf16 = DType.bf16.ZigType();188                const payload = try self.arena.allocator().alloc(Bf16, count);189                const value = Bf16.fromF32(@floatCast(fill_value));190                for (payload) |*slot| slot.* = value;191                break :blk self.constantOwned(owned, std.mem.sliceAsBytes(payload));192            },193            .f32 => blk: {194                const payload = try self.arena.allocator().alloc(f32, count);195                const value: f32 = @floatCast(fill_value);196                for (payload) |*slot| slot.* = value;197                break :blk self.constantOwned(owned, std.mem.sliceAsBytes(payload));198            },199            .f64 => blk: {200                const payload = try self.arena.allocator().alloc(f64, count);201                for (payload) |*slot| slot.* = fill_value;202                break :blk self.constantOwned(owned, std.mem.sliceAsBytes(payload));203            },204            else => error.NonFloatDType,205        };206    }207208    pub fn constantBytes(self: *Builder, ty: Type, payload: []const u8) !Value {209        if (payload.len != try ty.byteCount()) return error.PayloadLengthMismatch;210        return self.constantOwned(try self.copyType(ty), try self.arena.allocator().dupe(u8, payload));211    }212213    pub fn zeros(self: *Builder, ty: Type) !Value {214        const owned = try self.copyType(ty);215        const payload = try self.arena.allocator().alloc(u8, try owned.byteCount());216        @memset(payload, 0);217        return self.constantOwned(owned, payload);218    }219220    pub fn isStructuralZero(self: *Builder, id: Id) bool {221        const index: usize = @intCast(id.index);222        if (index >= self.operations.items.len) return false;223        return switch (self.operations.items[index].kind) {224            .constant => |constant| program_mod.isZeroPayload(constant.payload),225            .unary => |unary_op| switch (unary_op.op) {226                .neg => self.isStructuralZero(unary_op.input),227                else => false,228            },229            .binary => |binary_op| switch (binary_op.op) {230                .add, .sub => self.isStructuralZero(binary_op.lhs) and self.isStructuralZero(binary_op.rhs),231                .mul => self.isStructuralZero(binary_op.lhs) or self.isStructuralZero(binary_op.rhs),232                else => false,233            },234            .broadcast => |broadcast_op| self.isStructuralZero(broadcast_op.input),235            .broadcast_in_dim => |broadcast_op| self.isStructuralZero(broadcast_op.input),236            .reshape => |reshape_op| self.isStructuralZero(reshape_op.input),237            .transpose => |transpose_op| self.isStructuralZero(transpose_op.input),238            .gather => |gather_op| self.isStructuralZero(gather_op.input),239            .scatter_add => |scatter_add| self.isStructuralZero(scatter_add.input) and self.isStructuralZero(scatter_add.updates),240            .select => |select_op| self.isStructuralZero(select_op.on_true) and self.isStructuralZero(select_op.on_false),241            else => false,242        };243    }244245    pub fn iota(self: *Builder, dtype: DType, dims_struct: anytype, comptime axis: anytype) !Value {246        var buffer: [type_mod.dimCount(@TypeOf(dims_struct))]Dim = undefined;247        type_mod.fillDims(dims_struct, &buffer);248        return self.iotaDims(dtype, &buffer, comptime nameOf(axis));249    }250251    pub fn iotaDims(self: *Builder, dtype: DType, dims: []const Dim, axis_name: []const u8) !Value {252        try type_mod.validateAuthoredDims(dims);253        const index = type_mod.findDim(dims, axis_name) orelse return error.AxisNotFound;254        return self.iotaTyped(.{ .dtype = dtype, .dims = dims }, @intCast(index));255    }256257    fn iotaTyped(self: *Builder, ty: Type, axis: i64) !Value {258        try type_mod.validateAxes(ty.dims.len, &.{axis});259        return self.emit(try self.copyType(ty), .{ .iota = .{ .axis = axis } });260    }261262    pub fn unary(self: *Builder, op: Unary, operand_value: Value) !Value {263        try self.ensureValue(operand_value);264        return self.emit(try self.copyType(operand_value.ty), .{265            .unary = .{ .op = op, .input = operand_value.id },266        });267    }268269    pub fn binary(self: *Builder, op: Binary, lhs: Value, rhs: Value) !Value {270        try self.ensureValue(lhs);271        try self.ensureValue(rhs);272        try type_mod.sameType(lhs.ty, rhs.ty);273        return self.emit(try self.copyType(lhs.ty), .{274            .binary = .{ .op = op, .lhs = lhs.id, .rhs = rhs.id },275        });276    }277278    pub fn alignedBinary(self: *Builder, op: Binary, lhs: Value, rhs: Value) !Value {279        try self.ensureValue(lhs);280        try self.ensureValue(rhs);281        if (lhs.ty.dtype != rhs.ty.dtype) return error.DTypeMismatch;282        const target = try type_mod.unionDims(self.arena.allocator(), lhs.ty.dims, rhs.ty.dims);283        const left = try self.alignTo(lhs, target);284        const right = try self.alignTo(rhs, target);285        return self.binary(op, left, right);286    }287288    pub fn customCall(289        self: *Builder,290        target: []const u8,291        version: u32,292        operands: []const Value,293        result_ty: Type,294    ) !Value {295        if (target.len == 0) return error.InvalidCustomCallContract;296        if (version == 0) return error.InvalidCustomCallContract;297        if (operands.len > program_mod.max_custom_call_operands) return error.UnsupportedCustomCallArity;298        const arena = self.arena.allocator();299        const ids = try arena.alloc(program_mod.Id, operands.len);300        for (operands, ids) |operand, *id| {301            try self.ensureValue(operand);302            id.* = operand.id;303        }304        return self.emit(try self.copyType(result_ty), .{305            .custom_call = .{306                .target = try arena.dupe(u8, target),307                .version = version,308                .operands = ids,309            },310        });311    }312313    pub fn compare(self: *Builder, direction: CompareDirection, lhs: Value, rhs: Value) !Value {314        try self.ensureValue(lhs);315        try self.ensureValue(rhs);316        try type_mod.sameType(lhs.ty, rhs.ty);317        const pred_ty = try self.copyType(.{ .dtype = .i1, .dims = lhs.ty.dims });318        return self.emit(pred_ty, .{319            .compare = .{ .lhs = lhs.id, .rhs = rhs.id, .direction = direction },320        });321    }322323    pub fn alignedCompare(self: *Builder, direction: CompareDirection, lhs: Value, rhs: Value) !Value {324        try self.ensureValue(lhs);325        try self.ensureValue(rhs);326        if (lhs.ty.dtype != rhs.ty.dtype) return error.DTypeMismatch;327        const target = try type_mod.unionDims(self.arena.allocator(), lhs.ty.dims, rhs.ty.dims);328        const left = try self.alignTo(lhs, target);329        const right = try self.alignTo(rhs, target);330        return self.compare(direction, left, right);331    }332333    pub fn select(self: *Builder, pred: Value, on_true: Value, on_false: Value) !Value {334        try self.ensureValue(pred);335        try self.ensureValue(on_true);336        try self.ensureValue(on_false);337        const ty = try type_mod.select(self.arena.allocator(), pred.ty, on_true.ty, on_false.ty);338        return self.emit(ty, .{339            .select = .{ .pred = pred.id, .on_true = on_true.id, .on_false = on_false.id },340        });341    }342343    pub fn alignedSelect(self: *Builder, pred: Value, on_true: Value, on_false: Value) !Value {344        try self.ensureValue(pred);345        try self.ensureValue(on_true);346        try self.ensureValue(on_false);347        if (on_true.ty.dtype != on_false.ty.dtype) return error.DTypeMismatch;348        const arena = self.arena.allocator();349        var target = try type_mod.unionDims(arena, on_true.ty.dims, on_false.ty.dims);350        if (pred.ty.rank() != 0) {351            target = try type_mod.unionDims(arena, target, pred.ty.dims);352        }353        const chosen_true = try self.alignTo(on_true, target);354        const chosen_false = try self.alignTo(on_false, target);355        const chosen_pred = if (pred.ty.rank() == 0) pred else try self.alignTo(pred, target);356        return self.select(chosen_pred, chosen_true, chosen_false);357    }358359    pub fn broadcastAxes(self: *Builder, operand_value: Value, added: []const Dim) !Value {360        try self.ensureValue(operand_value);361        try type_mod.validateAuthoredDims(added);362        const target = try type_mod.appendDims(self.arena.allocator(), operand_value.ty.dims, added);363        return self.alignTo(operand_value, target);364    }365366    pub fn renameAxis(self: *Builder, operand_value: Value, old_name: []const u8, new_name: []const u8) !Value {367        try self.ensureValue(operand_value);368        try type_mod.validateAuthoredName(new_name);369        const new_dims = try type_mod.renamed(self.arena.allocator(), operand_value.ty.dims, old_name, new_name);370        return self.reshapeTo(operand_value, new_dims);371    }372373    pub fn splitAxis(self: *Builder, operand_value: Value, axis_name: []const u8, parts: []const Dim) !Value {374        try self.ensureValue(operand_value);375        try type_mod.validateAuthoredDims(parts);376        const new_dims = try type_mod.splitDims(self.arena.allocator(), operand_value.ty.dims, axis_name, parts);377        return self.reshapeTo(operand_value, new_dims);378    }379380    pub fn mergeAxes(self: *Builder, operand_value: Value, names: []const []const u8, merged_name: []const u8) !Value {381        try self.ensureValue(operand_value);382        try type_mod.validateAuthoredName(merged_name);383        const plan = try type_mod.mergeDims(self.arena.allocator(), operand_value.ty.dims, names, merged_name);384        var current = operand_value;385        if (plan.permutation) |permutation| {386            current = try self.transposeBy(current, permutation);387        }388        return self.reshapeTo(current, plan.result);389    }390391    pub fn alignTo(self: *Builder, operand_value: Value, target: []const Dim) !Value {392        if (type_mod.sameDims(operand_value.ty.dims, target)) return operand_value;393        const plan = try type_mod.alignment(self.arena.allocator(), operand_value.ty.dims, target);394        var current = operand_value;395        if (plan.permutation) |permutation| {396            current = try self.transposeBy(current, permutation);397        }398        if (type_mod.sameDims(current.ty.dims, target)) return current;399        return self.broadcastTo(current, target, plan.mapping.?);400    }401402    pub fn broadcastOp(self: *Builder, operand_value: Value, result_dims: []const Dim) !Value {403        try self.ensureValue(operand_value);404        const ty = try Type.init(self.arena.allocator(), operand_value.ty.dtype, result_dims);405        return self.emit(ty, .{406            .broadcast = .{407                .input = operand_value.id,408                .sizes = try type_mod.extents(self.arena.allocator(), ty.dims),409            },410        });411    }412413    pub fn broadcastTo(self: *Builder, operand_value: Value, result_dims: []const Dim, broadcast_dims: []const i64) !Value {414        try self.ensureValue(operand_value);415        const ty = try type_mod.broadcastInDim(self.arena.allocator(), operand_value.ty, result_dims, broadcast_dims);416        return self.emit(ty, .{417            .broadcast_in_dim = .{418                .input = operand_value.id,419                .broadcast_dims = try self.arena.allocator().dupe(i64, broadcast_dims),420            },421        });422    }423424    pub fn reshapeTo(self: *Builder, operand_value: Value, new_dims: []const Dim) !Value {425        try self.ensureValue(operand_value);426        const ty = try type_mod.reshaped(self.arena.allocator(), operand_value.ty, new_dims);427        return self.emit(ty, .{428            .reshape = .{429                .input = operand_value.id,430                .new_shape = try type_mod.extents(self.arena.allocator(), ty.dims),431            },432        });433    }434435    pub fn transposeBy(self: *Builder, operand_value: Value, permutation: []const i64) !Value {436        try self.ensureValue(operand_value);437        const result_dims = try type_mod.permuted(self.arena.allocator(), operand_value.ty.dims, permutation);438        const ty = try Type.init(self.arena.allocator(), operand_value.ty.dtype, result_dims);439        return self.emit(ty, .{440            .transpose = .{441                .input = operand_value.id,442                .permutation = try self.arena.allocator().dupe(i64, permutation),443            },444        });445    }446447    pub fn reduce(self: *Builder, operand_value: Value, init_value: Value, reducer: Reducer, dimensions: []const i64) !Value {448        try self.ensureValue(operand_value);449        try self.ensureValue(init_value);450        if (operand_value.ty.dtype != init_value.ty.dtype) return error.DTypeMismatch;451        if (init_value.ty.rank() != 0) return error.ReduceInitNotScalar;452        const result_dims = try type_mod.removeAxes(self.arena.allocator(), operand_value.ty.dims, dimensions);453        const ty = try Type.init(self.arena.allocator(), operand_value.ty.dtype, result_dims);454        return self.emit(ty, .{455            .reduce = .{456                .input = operand_value.id,457                .init = init_value.id,458                .reducer = reducer,459                .dimensions = try self.arena.allocator().dupe(i64, dimensions),460            },461        });462    }463464    pub fn reduceNamed(self: *Builder, operand_value: Value, reducer: Reducer, names: []const []const u8) !Value {465        try self.ensureValue(operand_value);466        const indices = try type_mod.axisIndices(self.arena.allocator(), operand_value.ty.dims, names);467        const init_value = try self.reducerInit(operand_value.ty.dtype, reducer);468        return self.reduce(operand_value, init_value, reducer, indices);469    }470471    pub fn reduceWith(self: *Builder, operand_value: Value, init_value: Value, reducer: Reducer, names: []const []const u8) !Value {472        try self.ensureValue(operand_value);473        const indices = try type_mod.axisIndices(self.arena.allocator(), operand_value.ty.dims, names);474        return self.reduce(operand_value, init_value, reducer, indices);475    }476477    pub fn gather(self: *Builder, operand_value: Value, indices_value: Value, comptime axis: anytype) !Value {478        return self.gatherNamed(operand_value, indices_value, comptime nameOf(axis));479    }480481    pub fn gatherNamed(self: *Builder, operand_value: Value, indices_value: Value, axis_name: []const u8) !Value {482        try self.ensureValue(operand_value);483        const axis = type_mod.findDim(operand_value.ty.dims, axis_name) orelse return error.AxisNotFound;484        return self.gatherOp(operand_value, indices_value, @intCast(axis));485    }486487    pub fn gatherOp(self: *Builder, operand_value: Value, indices_value: Value, axis: i64) !Value {488        try self.ensureValue(operand_value);489        try self.ensureValue(indices_value);490        const ty = try type_mod.gather(self.arena.allocator(), operand_value.ty, indices_value.ty, axis);491        return self.emit(ty, .{492            .gather = .{493                .input = operand_value.id,494                .indices = indices_value.id,495                .axis = axis,496            },497        });498    }499500    pub fn scatterAdd(self: *Builder, input_value: Value, indices_value: Value, updates_value: Value, comptime axis: anytype) !Value {501        return self.scatterAddNamed(input_value, indices_value, updates_value, comptime nameOf(axis));502    }503504    pub fn scatterAddNamed(self: *Builder, input_value: Value, indices_value: Value, updates_value: Value, axis_name: []const u8) !Value {505        try self.ensureValue(input_value);506        const axis = type_mod.findDim(input_value.ty.dims, axis_name) orelse return error.AxisNotFound;507        return self.scatterAddOp(input_value, indices_value, updates_value, @intCast(axis));508    }509510    pub fn scatterAddOp(self: *Builder, input_value: Value, indices_value: Value, updates_value: Value, axis: i64) !Value {511        try self.ensureValue(input_value);512        try self.ensureValue(indices_value);513        try self.ensureValue(updates_value);514        const ty = try type_mod.scatterAdd(self.arena.allocator(), input_value.ty, indices_value.ty, updates_value.ty, axis);515        return self.emit(ty, .{516            .scatter_add = .{517                .input = input_value.id,518                .indices = indices_value.id,519                .updates = updates_value.id,520                .axis = axis,521            },522        });523    }524525    pub fn sparseCrossEntropyLoss(self: *Builder, logits_value: Value, targets_value: Value, comptime axis: anytype) !Value {526        return self.sparseCrossEntropyLossNamed(logits_value, targets_value, comptime nameOf(axis));527    }528529    pub fn sparseCrossEntropyLossNamed(self: *Builder, logits_value: Value, targets_value: Value, axis_name: []const u8) !Value {530        try self.ensureValue(logits_value);531        const axis = type_mod.findDim(logits_value.ty.dims, axis_name) orelse return error.AxisNotFound;532        return self.sparseCrossEntropyOp(logits_value, targets_value, @intCast(axis));533    }534535    pub fn sparseCrossEntropyOp(self: *Builder, logits_value: Value, targets_value: Value, axis: i64) !Value {536        try self.ensureValue(logits_value);537        try self.ensureValue(targets_value);538        const ty = try type_mod.sparseCrossEntropy(self.arena.allocator(), logits_value.ty, targets_value.ty, axis);539        return self.emit(ty, .{540            .sparse_cross_entropy = .{541                .logits = logits_value.id,542                .targets = targets_value.id,543                .axis = axis,544            },545        });546    }547548    pub fn meanNamed(self: *Builder, operand_value: Value, names: []const []const u8) !Value {549        try self.ensureValue(operand_value);550        const indices = try type_mod.axisIndices(self.arena.allocator(), operand_value.ty.dims, names);551        const count = try type_mod.reducedExtentProduct(operand_value.ty.dims, indices);552        const init_value = try self.reducerInit(operand_value.ty.dtype, .sum);553        const summed = try self.reduce(operand_value, init_value, .sum, indices);554        const divisor = try self.fullFloat(Type.scalar(operand_value.ty.dtype), @floatFromInt(count));555        return self.alignedBinary(.div, summed, divisor);556    }557558    fn reducerInit(self: *Builder, dtype: DType, reducer: Reducer) !Value {559        const ty = Type.scalar(dtype);560        return switch (reducer) {561            .sum => self.zeros(ty),562            .max => self.extremum(ty, .lowest),563            .min => self.extremum(ty, .highest),564        };565    }566567    fn extremum(self: *Builder, ty: Type, comptime bound: enum { lowest, highest }) !Value {568        return switch (ty.dtype) {569            inline .f16, .f32, .f64 => |tag| blk: {570                const Element = tag.ZigType();571                const value: Element = if (bound == .lowest) -std.math.inf(Element) else std.math.inf(Element);572                break :blk self.constantBytes(ty, std.mem.asBytes(&value));573            },574            .bf16 => blk: {575                const Bf16 = DType.bf16.ZigType();576                const value = Bf16.fromF32(if (bound == .lowest) -std.math.inf(f32) else std.math.inf(f32));577                break :blk self.constantBytes(ty, std.mem.asBytes(&value));578            },579            inline .i32, .i64 => |tag| blk: {580                const Element = tag.ZigType();581                const value: Element = if (bound == .lowest) std.math.minInt(Element) else std.math.maxInt(Element);582                break :blk self.constantBytes(ty, std.mem.asBytes(&value));583            },584            else => error.DTypeMismatch,585        };586    }587588    pub fn contract(self: *Builder, lhs: Value, rhs: Value, names: []const []const u8) !Value {589        try self.ensureValue(lhs);590        try self.ensureValue(rhs);591        if (lhs.ty.dtype != rhs.ty.dtype) return error.ContractDTypeMismatch;592        const arena = self.arena.allocator();593        const plan = try type_mod.contraction(arena, lhs.ty.dims, rhs.ty.dims, names);594595        const batch_count = plan.lhs_batch.len;596        const contract_count = plan.lhs_contract.len;597        const lhs_free = lhs.ty.dims.len - batch_count - contract_count;598599        const lhs_order = try arena.alloc(Dim, lhs.ty.dims.len);600        for (plan.lhs_batch, 0..) |axis, index| {601            lhs_order[index] = lhs.ty.dims[@intCast(axis)];602        }603        fillFreeDims(lhs.ty.dims, plan.lhs_batch, plan.lhs_contract, lhs_order[batch_count .. batch_count + lhs_free]);604        for (plan.lhs_contract, 0..) |axis, index| {605            lhs_order[batch_count + lhs_free + index] = lhs.ty.dims[@intCast(axis)];606        }607608        const rhs_order = try arena.alloc(Dim, rhs.ty.dims.len);609        for (plan.rhs_batch, 0..) |axis, index| {610            rhs_order[index] = rhs.ty.dims[@intCast(axis)];611        }612        for (plan.rhs_contract, 0..) |axis, index| {613            rhs_order[batch_count + index] = rhs.ty.dims[@intCast(axis)];614        }615        fillFreeDims(rhs.ty.dims, plan.rhs_batch, plan.rhs_contract, rhs_order[batch_count + contract_count ..]);616617        const lhs_canonical = try self.alignTo(lhs, lhs_order);618        const rhs_canonical = try self.alignTo(rhs, rhs_order);619620        const lhs_batch = try arena.alloc(i64, batch_count);621        const rhs_batch = try arena.alloc(i64, batch_count);622        for (0..batch_count) |index| {623            lhs_batch[index] = @intCast(index);624            rhs_batch[index] = @intCast(index);625        }626        const lhs_contract = try arena.alloc(i64, contract_count);627        const rhs_contract = try arena.alloc(i64, contract_count);628        for (0..contract_count) |index| {629            lhs_contract[index] = @intCast(batch_count + lhs_free + index);630            rhs_contract[index] = @intCast(batch_count + index);631        }632633        const ty = try Type.init(arena, lhs.ty.dtype, plan.result);634        return self.emit(ty, .{635            .dot_general = .{636                .lhs = lhs_canonical.id,637                .rhs = rhs_canonical.id,638                .lhs_contract = lhs_contract,639                .rhs_contract = rhs_contract,640                .lhs_batch = lhs_batch,641                .rhs_batch = rhs_batch,642            },643        });644    }645646    pub fn dotGeneralOp(647        self: *Builder,648        lhs: Value,649        rhs: Value,650        lhs_contract: []const i64,651        rhs_contract: []const i64,652        lhs_batch: []const i64,653        rhs_batch: []const i64,654    ) !Value {655        try self.ensureValue(lhs);656        try self.ensureValue(rhs);657        if (lhs.ty.dtype != rhs.ty.dtype) return error.ContractDTypeMismatch;658        const arena = self.arena.allocator();659        const result_dims = try type_mod.dotGeneralDims(660            arena,661            lhs.ty.dims,662            rhs.ty.dims,663            lhs_contract,664            rhs_contract,665            lhs_batch,666            rhs_batch,667        );668        const ty = try Type.init(arena, lhs.ty.dtype, result_dims);669        return self.emit(ty, .{670            .dot_general = .{671                .lhs = lhs.id,672                .rhs = rhs.id,673                .lhs_contract = try arena.dupe(i64, lhs_contract),674                .rhs_contract = try arena.dupe(i64, rhs_contract),675                .lhs_batch = try arena.dupe(i64, lhs_batch),676                .rhs_batch = try arena.dupe(i64, rhs_batch),677            },678        });679    }680681    pub fn scan(self: *Builder, spec: anytype) !ScanResult(@TypeOf(spec.init)) {682        const init_count = comptime scanValueCount(@TypeOf(spec.init));683        var inits: [init_count]Value = undefined;684        scanValues(spec.init, &inits);685686        var scope = try self.scanScope(spec.length, &inits);687        errdefer scope.abort();688689        const carry = scanCarry(@TypeOf(spec.init), &scope);690        const next = try callScanBody(spec.body, scope.body(), carry);691        var next_values: [init_count]Value = undefined;692        scanValues(next, &next_values);693694        const walked = try scope.finish(&next_values);695        return scanResult(self, @TypeOf(spec.init), walked);696    }697698    pub fn scanScope(self: *Builder, length: i64, inits: []const Value) !ScanScope {699        if (length < 0) return error.ScanLengthNegative;700        if (inits.len == 0) return error.ScanWithoutCarries;701        if (inits.len > program_mod.max_scan_carries) return error.ScanArityUnsupported;702        const init_ids = try self.arena.allocator().alloc(Id, inits.len);703        for (inits, init_ids) |value, *slot| {704            try self.ensureValue(value);705            slot.* = value.id;706        }707        var child = try Builder.init(self.allocator, "scan_body");708        errdefer child.deinit();709        for (inits) |value| {710            _ = try child.inputTyped(value.ty);711        }712        return .{713            .parent = self,714            .child = child,715            .length = length,716            .init_ids = init_ids,717        };718    }719720    pub fn projection(self: *Builder, source: Value, index: usize) !Value {721        try self.ensureValue(source);722        const source_op = self.operations.items[source.id.index];723        const scan_op = switch (source_op.kind) {724            .scan => |scan_op| scan_op,725            else => return error.ProjectionSourceNotScan,726        };727        if (index == 0 or index >= scan_op.inits.len) return error.ProjectionIndexInvalid;728        const ty = try self.copyType(scan_op.body.typeOf(scan_op.body.outputs[index]));729        return self.emit(ty, .{ .projection = .{ .source = source.id, .index = index } });730    }731732    pub fn emitScan(self: *Builder, length: i64, inits: []const Value, body: *const program_mod.Subgraph) !Value {733        const init_ids = try self.arena.allocator().alloc(Id, inits.len);734        for (inits, init_ids) |value, *slot| {735            try self.ensureValue(value);736            slot.* = value.id;737        }738        return self.emitScanIds(length, init_ids, body);739    }740741    fn emitScanIds(self: *Builder, length: i64, inits: []const Id, body: *const program_mod.Subgraph) !Value {742        if (length < 0) return error.ScanLengthNegative;743        if (inits.len == 0) return error.ScanWithoutCarries;744        if (inits.len > program_mod.max_scan_carries) return error.ScanArityUnsupported;745        if (body.parameters.len != inits.len) return error.ScanArityMismatch;746        if (body.outputs.len != inits.len) return error.ScanArityMismatch;747        for (inits, 0..) |init_id, index| {748            const init_ty = self.values.items[init_id.index];749            try type_mod.sameType(init_ty, body.typeOf(body.parameters[index]));750            try type_mod.sameType(init_ty, body.typeOf(body.outputs[index]));751        }752        const owned_body = try program_mod.cloneSubgraph(self.arena.allocator(), body);753        const result_ty = try self.copyType(self.values.items[inits[0].index]);754        return self.emit(result_ty, .{ .scan = .{755            .length = length,756            .inits = inits,757            .body = owned_body,758        } });759    }760761    pub fn finish(self: *Builder, outputs: []const Value) !Program {762        if (self.finished) return error.BuilderFinished;763        const owned_outputs = try self.arena.allocator().alloc(Id, outputs.len);764        for (outputs, 0..) |output, index| {765            try self.ensureValue(output);766            owned_outputs[index] = output.id;767        }768        self.finished = true;769        return .{770            .arena = self.arena,771            .name = self.name,772            .values = self.values.items,773            .operations = self.operations.items,774            .parameters = self.parameters_list.items,775            .outputs = owned_outputs,776        };777    }778779    fn constantOwned(self: *Builder, ty: Type, payload: []const u8) !Value {780        return self.emit(ty, .{ .constant = .{ .payload = payload } });781    }782783    fn emit(self: *Builder, ty: Type, kind: program_mod.Kind) !Value {784        const id = try self.append(ty, kind);785        return self.makeValue(id);786    }787788    fn append(self: *Builder, result: Type, kind: program_mod.Kind) !Id {789        if (self.finished) return error.BuilderFinished;790        const id = Id{ .index = @intCast(self.values.items.len) };791        try self.values.append(self.arena.allocator(), result);792        try self.operations.append(self.arena.allocator(), .{793            .id = id,794            .result = result,795            .kind = kind,796        });797        return id;798    }799800    fn makeValue(self: *Builder, id: Id) Value {801        return .{802            .builder = self,803            .id = id,804            .ty = self.values.items[id.index],805        };806    }807808    fn ensureValue(self: *Builder, candidate: Value) !void {809        if (self.finished) return error.BuilderFinished;810        if (candidate.builder != self) return error.CrossBuilderValue;811    }812813    fn copyType(self: *Builder, ty: Type) !Type {814        return Type.init(self.arena.allocator(), ty.dtype, ty.dims);815    }816};817818fn fillFreeDims(dims: []const Dim, batch: []const i64, contracted: []const i64, out: []Dim) void {819    var index: usize = 0;820    for (dims, 0..) |dim, position| {821        if (containsIndex(batch, position) or containsIndex(contracted, position)) continue;822        out[index] = dim;823        index += 1;824    }825}826827fn containsIndex(axes: []const i64, position: usize) bool {828    for (axes) |axis| {829        if (axis == @as(i64, @intCast(position))) return true;830    }831    return false;832}833834pub fn nameOf(comptime name_value: anytype) []const u8 {835    return switch (@typeInfo(@TypeOf(name_value))) {836        .enum_literal => @tagName(name_value),837        .pointer, .array => name_value[0..],838        else => @compileError("tensor axis names must be enum literals or string literals"),839    };840}841842fn ScanResult(comptime Init: type) type {843    if (comptime Init == Value) return Value;844    return switch (@typeInfo(Init)) {845        .array => |array| blk: {846            if (array.child != Value) @compileError("tensor scan array init values must be tensor Values");847            break :blk [array.len]Value;848        },849        .@"struct" => |info| blk: {850            if (info.field_names.len == 0) @compileError("tensor scan init struct must contain at least one carry");851            inline for (info.field_types) |field_type| {852                if (field_type != Value) @compileError("tensor scan struct init fields must be tensor Values");853            }854            break :blk Init;855        },856        else => @compileError("tensor scan init must be a Value, array of Values, tuple of Values, or struct of Values"),857    };858}859860fn scanValueCount(comptime Init: type) usize {861    if (comptime Init == Value) return 1;862    return switch (@typeInfo(Init)) {863        .array => |array| blk: {864            if (array.child != Value) @compileError("tensor scan array init values must be tensor Values");865            break :blk array.len;866        },867        .@"struct" => |info| blk: {868            if (info.field_names.len == 0) @compileError("tensor scan init struct must contain at least one carry");869            inline for (info.field_types) |field_type| {870                if (field_type != Value) @compileError("tensor scan struct init fields must be tensor Values");871            }872            break :blk info.field_names.len;873        },874        else => @compileError("tensor scan init must be a Value, array of Values, tuple of Values, or struct of Values"),875    };876}877878fn scanValues(source: anytype, values: *[scanValueCount(@TypeOf(source))]Value) void {879    const Source = @TypeOf(source);880    if (comptime Source == Value) {881        values[0] = source;882        return;883    }884    switch (@typeInfo(Source)) {885        .array => for (source, 0..) |value, index| {886            values[index] = value;887        },888        .@"struct" => |info| inline for (info.field_names, 0..) |field_name, index| {889            values[index] = @field(source, field_name);890        },891        else => unreachable,892    }893}894895fn scanCarry(comptime Init: type, scope: *ScanScope) ScanResult(Init) {896    if (comptime Init == Value) return scope.carry(0);897    return switch (@typeInfo(Init)) {898        .array => |array| blk: {899            var result: [array.len]Value = undefined;900            inline for (0..array.len) |index| {901                result[index] = scope.carry(index);902            }903            break :blk result;904        },905        .@"struct" => |info| blk: {906            var result: Init = undefined;907            inline for (info.field_names, 0..) |field_name, index| {908                @field(result, field_name) = scope.carry(index);909            }910            break :blk result;911        },912        else => unreachable,913    };914}915916fn scanResult(builder: *Builder, comptime Init: type, walked: Value) !ScanResult(Init) {917    if (comptime Init == Value) return walked;918    return switch (@typeInfo(Init)) {919        .array => |array| blk: {920            var result: [array.len]Value = undefined;921            inline for (0..array.len) |index| {922                result[index] = if (index == 0) walked else try builder.projection(walked, index);923            }924            break :blk result;925        },926        .@"struct" => |info| blk: {927            var result: Init = undefined;928            inline for (info.field_names, 0..) |field_name, index| {929                @field(result, field_name) = if (index == 0) walked else try builder.projection(walked, index);930            }931            break :blk result;932        },933        else => unreachable,934    };935}936937fn callScanBody(comptime body: anytype, builder: *Builder, carry: anytype) !ScanResult(@TypeOf(carry)) {938    if (comptime @typeInfo(@TypeOf(body)) != .@"fn") @compileError("tensor scan body must be a function value");939    return body(builder, carry);940}941942test "indexing operations trace with named axis shapes" {943    var builder = try Builder.init(std.testing.allocator, "trace_indexing");944    defer builder.deinit();945946    const table = try builder.input(.f32, .{ .vocab = 32, .channel = 8 });947    const ids = try builder.input(.i32, .{ .token = 5 });948    const gathered = try table.gather(ids, .vocab);949    try std.testing.expectEqual(@as(usize, 2), gathered.ty.rank());950    try std.testing.expectEqualStrings("token", gathered.ty.dims[0].name);951    try std.testing.expectEqualStrings("channel", gathered.ty.dims[1].name);952953    const zero_table = try builder.full(.f32, .{ .vocab = 32, .channel = 8 }, 0.0);954    const accumulated = try zero_table.scatterAdd(ids, gathered, .vocab);955    try std.testing.expect(type_mod.sameDims(table.ty.dims, accumulated.ty.dims));956957    var program = try builder.finish(&.{accumulated});958    defer program.deinit();959960    try std.testing.expect(program.operation(gathered.id).kind == .gather);961    try std.testing.expect(program.operation(accumulated.id).kind == .scatter_add);962}963964test "scan scope traces a carry loop into one scan operation" {965    var builder = try Builder.init(std.testing.allocator, "scan_trace");966    defer builder.deinit();967968    const x0 = try builder.input(.f32, .{ .lane = 4 });969    const acc0 = try builder.full(.f32, .{ .lane = 4 }, 0.0);970971    var scope = try builder.scanScope(5, &.{ x0, acc0 });972    errdefer scope.abort();973    const half = try scope.body().scalar(.f32, 0.5);974    const next_x = try (try scope.carry(0).mul(scope.carry(0))).mul(half);975    const next_acc = try scope.carry(1).add(next_x);976    const walked = try scope.finish(&.{ next_x, next_acc });977    const final_acc = try builder.projection(walked, 1);978979    var program = try builder.finish(&.{ walked, final_acc });980    defer program.deinit();981982    const scan_op = program.operation(walked.id);983    try std.testing.expectEqual(@as(i64, 5), scan_op.kind.scan.length);984    try std.testing.expectEqual(@as(usize, 2), scan_op.kind.scan.inits.len);985    try std.testing.expectEqual(@as(usize, 2), scan_op.kind.scan.body.parameters.len);986    try std.testing.expectEqual(@as(usize, 2), scan_op.kind.scan.body.outputs.len);987    try std.testing.expectEqualStrings("lane", program.typeOf(walked.id).dims[0].name);988    try std.testing.expectEqual(@as(i64, 4), program.typeOf(final_acc.id).dims[0].extent);989990    const projection_op = program.operation(final_acc.id);991    try std.testing.expectEqual(walked.id, projection_op.kind.projection.source);992    try std.testing.expectEqual(@as(usize, 1), projection_op.kind.projection.index);993}994995fn namedScanStep(scan_builder: *Builder, carry: anytype) !@TypeOf(carry) {996    const half = try scan_builder.scalar(.f32, 0.5);997    const next_x = try (try carry.x.mul(carry.x)).mul(half);998    return .{999        .x = next_x,1000        .acc = try carry.acc.add(next_x),1001    };1002}10031004test "scan definition returns named carries" {1005    var builder = try Builder.init(std.testing.allocator, "scan_named");1006    defer builder.deinit();10071008    const x0 = try builder.input(.f32, .{ .lane = 4 });1009    const acc0 = try builder.full(.f32, .{ .lane = 4 }, 0.0);10101011    const walked = try builder.scan(.{1012        .length = 5,1013        .init = .{ .x = x0, .acc = acc0 },1014        .body = namedScanStep,1015    });10161017    var program = try builder.finish(&.{ walked.x, walked.acc });1018    defer program.deinit();10191020    const scan_op = program.operation(walked.x.id);1021    try std.testing.expectEqual(@as(i64, 5), scan_op.kind.scan.length);1022    try std.testing.expectEqual(@as(usize, 2), scan_op.kind.scan.inits.len);1023    try std.testing.expectEqual(@as(usize, 2), scan_op.kind.scan.body.outputs.len);10241025    const projection_op = program.operation(walked.acc.id);1026    try std.testing.expectEqual(walked.x.id, projection_op.kind.projection.source);1027    try std.testing.expectEqual(@as(usize, 1), projection_op.kind.projection.index);1028}10291030fn singleScanStep(_: *Builder, carry: Value) !Value {1031    return carry.add(carry);1032}10331034test "scan definition returns a single carry value" {1035    var builder = try Builder.init(std.testing.allocator, "scan_single");1036    defer builder.deinit();10371038    const x0 = try builder.input(.f32, .{ .lane = 4 });10391040    const walked = try builder.scan(.{1041        .length = 3,1042        .init = x0,1043        .body = singleScanStep,1044    });10451046    var program = try builder.finish(&.{walked});1047    defer program.deinit();10481049    const scan_op = program.operation(walked.id);1050    try std.testing.expectEqual(@as(i64, 3), scan_op.kind.scan.length);1051    try std.testing.expectEqual(@as(usize, 1), scan_op.kind.scan.inits.len);1052    try std.testing.expectEqual(@as(usize, 1), scan_op.kind.scan.body.outputs.len);1053}10541055fn arrayScanStep(_: *Builder, carry: [2]Value) ![2]Value {1056    const next = try carry[0].add(carry[0]);1057    return .{ next, try carry[1].add(next) };1058}10591060fn tupleScanStep(_: *Builder, carry: anytype) !@TypeOf(carry) {1061    const next = try carry[0].add(carry[0]);1062    return .{ next, try carry[1].add(next) };1063}10641065test "scan definition returns array and tuple carries" {1066    var array_builder = try Builder.init(std.testing.allocator, "scan_array");1067    defer array_builder.deinit();10681069    const array_x0 = try array_builder.input(.f32, .{ .lane = 2 });1070    const array_acc0 = try array_builder.full(.f32, .{ .lane = 2 }, 0.0);1071    const array_walked = try array_builder.scan(.{1072        .length = 2,1073        .init = [2]Value{ array_x0, array_acc0 },1074        .body = arrayScanStep,1075    });1076    var array_program = try array_builder.finish(&.{ array_walked[0], array_walked[1] });1077    defer array_program.deinit();10781079    const array_projection = array_program.operation(array_walked[1].id);1080    try std.testing.expectEqual(array_walked[0].id, array_projection.kind.projection.source);1081    try std.testing.expectEqual(@as(usize, 1), array_projection.kind.projection.index);10821083    var tuple_builder = try Builder.init(std.testing.allocator, "scan_tuple");1084    defer tuple_builder.deinit();10851086    const tuple_x0 = try tuple_builder.input(.f32, .{ .lane = 2 });1087    const tuple_acc0 = try tuple_builder.full(.f32, .{ .lane = 2 }, 0.0);1088    const tuple_walked = try tuple_builder.scan(.{1089        .length = 2,1090        .init = .{ tuple_x0, tuple_acc0 },1091        .body = tupleScanStep,1092    });1093    var tuple_program = try tuple_builder.finish(&.{ tuple_walked[0], tuple_walked[1] });1094    defer tuple_program.deinit();10951096    const tuple_projection = tuple_program.operation(tuple_walked[1].id);1097    try std.testing.expectEqual(tuple_walked[0].id, tuple_projection.kind.projection.source);1098    try std.testing.expectEqual(@as(usize, 1), tuple_projection.kind.projection.index);1099}11001101test "scan scope rejects misuse" {1102    var builder = try Builder.init(std.testing.allocator, "scan_misuse");1103    defer builder.deinit();11041105    const x0 = try builder.input(.f32, .{ .lane = 2 });11061107    try std.testing.expectError(error.ScanLengthNegative, builder.scanScope(-1, &.{x0}));1108    try std.testing.expectError(error.ScanWithoutCarries, builder.scanScope(3, &.{}));11091110    var scope = try builder.scanScope(3, &.{x0});1111    defer scope.abort();1112    const stranger = try builder.full(.f32, .{ .lane = 2 }, 1.0);1113    try std.testing.expectError(error.CrossBuilderValue, scope.finish(&.{stranger}));11141115    var second = try builder.scanScope(3, &.{x0});1116    defer second.abort();1117    const wide = try second.body().full(.f32, .{ .lane = 3 }, 0.0);1118    try std.testing.expectError(error.ShapeMismatch, second.finish(&.{wide}));11191120    var third = try builder.scanScope(3, &.{x0});1121    defer third.abort();1122    try std.testing.expectError(error.ScanArityMismatch, third.finish(&.{ third.carry(0), third.carry(0) }));11231124    const not_scan = try builder.full(.f32, .{ .lane = 2 }, 2.0);1125    try std.testing.expectError(error.ProjectionSourceNotScan, builder.projection(not_scan, 1));11261127    var fourth = try builder.scanScope(2, &.{x0});1128    const bumped = try fourth.carry(0).add(try fourth.body().full(.f32, .{ .lane = 2 }, 1.0));1129    const walked = try fourth.finish(&.{bumped});1130    try std.testing.expectError(error.ProjectionIndexInvalid, builder.projection(walked, 0));1131    try std.testing.expectError(error.ProjectionIndexInvalid, builder.projection(walked, 1));1132}11331134test "scan operations replay through the structural graph copy" {1135    var builder = try Builder.init(std.testing.allocator, "scan_replay_source");1136    defer builder.deinit();11371138    const x0 = try builder.input(.f32, .{ .lane = 3 });1139    const gain = try builder.input(.f32, .{ .lane = 3 });11401141    var scope = try builder.scanScope(4, &.{ x0, gain });1142    const next = try scope.carry(0).mul(scope.carry(1));1143    const walked = try scope.finish(&.{ next, scope.carry(1) });1144    const kept_gain = try builder.projection(walked, 1);1145    var source = try builder.finish(&.{ walked, kept_gain });1146    defer source.deinit();11471148    var replay_builder = try Builder.init(std.testing.allocator, "scan_replay_source");1149    errdefer replay_builder.deinit();1150    var values = std.ArrayListUnmanaged(Value).empty;1151    defer values.deinit(std.testing.allocator);1152    for (source.operations) |*op| {1153        var buffer: [program_mod.max_operation_operands]Value = undefined;1154        var count: usize = 0;1155        switch (op.kind) {1156            .parameter => {},1157            .scan => |scan| {1158                for (scan.inits) |init_id| {1159                    buffer[count] = values.items[init_id.index];1160                    count += 1;1161                }1162            },1163            .projection => |projection_op| {1164                buffer[0] = values.items[projection_op.source.index];1165                count = 1;1166            },1167            else => unreachable,1168        }1169        try values.append(std.testing.allocator, try replay_builder.operation(op, buffer[0..count]));1170    }1171    var outputs: [2]Value = .{ values.items[source.outputs[0].index], values.items[source.outputs[1].index] };1172    var replayed = try replay_builder.finish(&outputs);1173    defer replayed.deinit();11741175    try std.testing.expectEqual(source.fingerprint(), replayed.fingerprint());1176}11771178test "scan fingerprints see body changes" {1179    var first_builder = try Builder.init(std.testing.allocator, "scan_fp");1180    defer first_builder.deinit();1181    const first_x = try first_builder.input(.f32, .{ .lane = 2 });1182    var first_scope = try first_builder.scanScope(3, &.{first_x});1183    const first_next = try first_scope.carry(0).add(try first_scope.body().full(.f32, .{ .lane = 2 }, 1.0));1184    const first_out = try first_scope.finish(&.{first_next});1185    var first = try first_builder.finish(&.{first_out});1186    defer first.deinit();11871188    var second_builder = try Builder.init(std.testing.allocator, "scan_fp");1189    defer second_builder.deinit();1190    const second_x = try second_builder.input(.f32, .{ .lane = 2 });1191    var second_scope = try second_builder.scanScope(3, &.{second_x});1192    const second_next = try second_scope.carry(0).add(try second_scope.body().full(.f32, .{ .lane = 2 }, 2.0));1193    const second_out = try second_scope.finish(&.{second_next});1194    var second = try second_builder.finish(&.{second_out});1195    defer second.deinit();11961197    try std.testing.expect(first.fingerprint() != second.fingerprint());1198}11991200test "aligned binary broadcasts missing axes by name" {1201    var builder = try Builder.init(std.testing.allocator, "align_binary");1202    defer builder.deinit();12031204    const xs = try builder.input(.f32, .{ .point = 4 });1205    const noise = try builder.input(.f32, .{ .point = 4, .sample = 8 });1206    const sum = try xs.add(noise);12071208    try std.testing.expectEqual(@as(usize, 2), sum.ty.rank());1209    try std.testing.expectEqualStrings("point", sum.ty.dims[0].name);1210    try std.testing.expectEqualStrings("sample", sum.ty.dims[1].name);12111212    const flipped = try noise.add(xs);1213    try std.testing.expectEqualStrings("point", flipped.ty.dims[0].name);12141215    const scalar_gain = try builder.scalar(.f32, 2.0);1216    const scaled = try noise.mul(scalar_gain);1217    try std.testing.expect(type_mod.sameDims(noise.ty.dims, scaled.ty.dims));12181219    const wrong = try builder.input(.f32, .{ .point = 5 });1220    try std.testing.expectError(error.AxisExtentMismatch, xs.add(wrong));1221}12221223test "aligned binary transposes shared axes into agreement" {1224    var builder = try Builder.init(std.testing.allocator, "align_transpose");1225    defer builder.deinit();12261227    const row_major = try builder.input(.f32, .{ .row = 2, .col = 3 });1228    const col_major = try builder.input(.f32, .{ .col = 3, .row = 2 });1229    const sum = try row_major.add(col_major);12301231    try std.testing.expectEqualStrings("row", sum.ty.dims[0].name);1232    try std.testing.expectEqualStrings("col", sum.ty.dims[1].name);12331234    const transpose_op = builder.operations.items[sum.id.index - 1];1235    try std.testing.expect(transpose_op.kind == .transpose);1236}12371238test "named reductions use canonical inits" {1239    var builder = try Builder.init(std.testing.allocator, "named_reduce");1240    defer builder.deinit();12411242    const x = try builder.input(.f32, .{ .point = 4, .sample = 8 });12431244    const summed = try x.sum(.sample);1245    try std.testing.expectEqual(@as(usize, 1), summed.ty.rank());1246    try std.testing.expectEqualStrings("point", summed.ty.dims[0].name);12471248    const biggest = try x.max(.{ .point, .sample });1249    try std.testing.expectEqual(@as(usize, 0), biggest.ty.rank());12501251    const smallest = try x.min(.point);1252    try std.testing.expectEqualStrings("sample", smallest.ty.dims[0].name);12531254    const averaged = try x.mean(.sample);1255    try std.testing.expectEqualStrings("point", averaged.ty.dims[0].name);12561257    try std.testing.expectError(error.AxisNotFound, x.sum(.missing));1258}12591260test "contract derives matmul attention and batched shapes" {1261    var builder = try Builder.init(std.testing.allocator, "named_contract");1262    defer builder.deinit();12631264    const q = try builder.input(.f32, .{ .pos = 8, .head = 16 });1265    const k = try builder.input(.f32, .{ .ctx = 6, .head = 16 });1266    const scores = try q.contract(k, .head);1267    try std.testing.expectEqualStrings("pos", scores.ty.dims[0].name);1268    try std.testing.expectEqualStrings("ctx", scores.ty.dims[1].name);12691270    const v = try builder.input(.f32, .{ .ctx = 6, .val = 32 });1271    const out = try scores.contract(v, .ctx);1272    try std.testing.expectEqualStrings("pos", out.ty.dims[0].name);1273    try std.testing.expectEqualStrings("val", out.ty.dims[1].name);12741275    const lhs = try builder.input(.f32, .{ .walk = 5, .m = 2, .k = 3 });1276    const rhs = try builder.input(.f32, .{ .k = 3, .walk = 5, .n = 4 });1277    const batched = try lhs.contract(rhs, .k);1278    try std.testing.expectEqualStrings("walk", batched.ty.dims[0].name);1279    try std.testing.expectEqualStrings("m", batched.ty.dims[1].name);1280    try std.testing.expectEqualStrings("n", batched.ty.dims[2].name);1281}12821283test "structure ops rename split merge and broadcast by name" {1284    var builder = try Builder.init(std.testing.allocator, "named_structure");1285    defer builder.deinit();12861287    const x = try builder.input(.f32, .{ .pixels = 12 });12881289    const image = try x.split(.pixels, .{ .row = 3, .col = 4 });1290    try std.testing.expectEqualStrings("row", image.ty.dims[0].name);1291    try std.testing.expectEqualStrings("col", image.ty.dims[1].name);12921293    const relabeled = try image.rename(.row, .line);1294    try std.testing.expectEqualStrings("line", relabeled.ty.dims[0].name);12951296    const flat = try relabeled.merge(.{ .line, .col }, .pixels);1297    try std.testing.expectEqual(@as(usize, 1), flat.ty.rank());1298    try std.testing.expectEqual(@as(i64, 12), flat.ty.dims[0].extent);12991300    const widened = try flat.broadcast(.{ .sample = 2 });1301    try std.testing.expectEqualStrings("pixels", widened.ty.dims[0].name);1302    try std.testing.expectEqualStrings("sample", widened.ty.dims[1].name);13031304    try std.testing.expectError(error.AxisNameReserved, flat.broadcast(.{ .@"#batch" = 2 }));1305}13061307test "aligned select broadcasts predicate and branches" {1308    var builder = try Builder.init(std.testing.allocator, "named_select");1309    defer builder.deinit();13101311    const pred = try builder.input(.i1, .{ .point = 4 });1312    const wide = try builder.input(.f32, .{ .point = 4, .sample = 8 });1313    const fallback = try builder.scalar(.f32, 0.0);13141315    const chosen = try pred.select(wide, fallback);1316    try std.testing.expectEqual(@as(usize, 2), chosen.ty.rank());1317    try std.testing.expectEqualStrings("point", chosen.ty.dims[0].name);13181319    const scalar_pred = try builder.scalar(.i1, true);1320    const scalar_chosen = try scalar_pred.select(wide, fallback);1321    try std.testing.expectEqual(@as(usize, 2), scalar_chosen.ty.rank());1322}