tiny.accy.tensor.autodiff

1const std = @import("std");2const emit = @import("emit.zig");3const hook = @import("hook.zig");4const interpret = @import("interpret/root.zig");5const program_mod = @import("program.zig");6const trace = @import("trace/root.zig");7const transform = @import("transform.zig");8const unroll = @import("unroll.zig");9const types = @import("type/root.zig");1011pub const LinearizeOptions = struct {12    wrt: []const usize,13};1415pub const JvpOptions = LinearizeOptions;1617pub const Dual = struct {18    primal: trace.Value,19    tangent: trace.Value,20};2122const ResultMode = enum {23    linearization,24    downstream,25};2627pub const Linearization = struct {28    program: program_mod.Program,29    differentiated_parameters: []const usize,30    primal_parameter_count: usize,31    tangent_parameter_count: usize,32    primal_output_count: usize,33    tangent_output_count: usize,3435    pub fn deinit(self: *Linearization) void {36        self.program.deinit();37    }3839    pub fn primalOutputs(self: *const Linearization) []const program_mod.Id {40        return self.program.outputs[0..self.primal_output_count];41    }4243    pub fn tangentOutputs(self: *const Linearization) []const program_mod.Id {44        return self.program.outputs[self.primal_output_count .. self.primal_output_count + self.tangent_output_count];45    }46};4748pub fn semantics(49    source: *const program_mod.Program,50    next: anytype,51    options: LinearizeOptions,52) Semantics(@TypeOf(next), .linearization) {53    return .{54        .source = source,55        .next = next,56        .options = options,57    };58}5960pub fn jvpSemantics(61    source: *const program_mod.Program,62    next: anytype,63    options: JvpOptions,64) Semantics(@TypeOf(next), .downstream) {65    return .{66        .source = source,67        .next = next,68        .options = options,69    };70}7172pub fn jvpSemanticsWith(73    source: *const program_mod.Program,74    next: anytype,75    options: JvpOptions,76    hooks: anytype,77) @TypeOf(hook.attach("jvp", hooks, jvpSemantics(source, hook.attach("linearize", hooks, next), options))) {78    const generated = hook.attach("linearize", hooks, next);79    const linear = jvpSemantics(source, generated, options);80    return hook.attach("jvp", hooks, linear);81}8283pub fn Semantics(comptime Next: type, comptime result_mode: ResultMode) type {84    return struct {85        source: *const program_mod.Program,86        next: Next,87        options: LinearizeOptions,88        parameter_values: ?[]Value = null,8990        pub const Value: type = Dual;91        pub const Result: type = if (result_mode == .linearization) Linearization else Next.Result;9293        pub fn operation(self: *@This(), step: *interpret.Step(Value)) !Value {94            var buffer: [program_mod.max_operation_operands]Value = undefined;95            return self.bind(step.op, interpret.arguments(Value, step.op, step.values, &buffer));96        }9798        pub fn bind(self: *@This(), op: *const program_mod.Operation, args: []const Value) !Value {99            return switch (op.kind) {100                .parameter => |parameter_info| self.parameterValue(parameter_info),101                .constant => self.constantValue(op),102                .iota => self.iotaValue(op),103                .unary => |unary_info| self.unaryValue(op, unary_info, args[0]),104                .binary => |binary_info| self.binaryValue(op, binary_info, args[0], args[1]),105                .broadcast => self.broadcastValue(op, args[0]),106                .broadcast_in_dim => |broadcast_info| self.broadcastInDimValue(op, broadcast_info, args[0]),107                .reshape => self.reshapeValue(op, args[0]),108                .transpose => |transpose_info| self.transposeValue(op, transpose_info, args[0]),109                .reduce => |reduce_info| self.reduceValue(op, reduce_info, args[0], args[1]),110                .gather => |gather_info| self.gatherValue(op, gather_info, args[0], args[1]),111                .scatter_add => |scatter_add| self.scatterAddValue(op, scatter_add, args[0], args[1], args[2]),112                .sparse_cross_entropy => |sparse_cross_entropy| self.sparseCrossEntropyValue(op, sparse_cross_entropy, args[0], args[1]),113                .dot_general => |dot_info| self.dotGeneralValue(op, dot_info, args[0], args[1]),114                .compare => self.compareValue(op, args[0], args[1]),115                .select => self.selectValue(op, args[0], args[1], args[2]),116                .custom_call => error.CustomCallRequiresJvpContract,117                .scan, .projection => error.ScanLinearizationUnsupported,118            };119        }120121        pub fn finish(self: *@This(), outputs: []const Value) !Result {122            const values = try self.builderHandle().arena.allocator().alloc(trace.Value, outputs.len * 2);123            for (outputs, 0..) |output, index| {124                values[index] = output.primal;125                values[index + outputs.len] = output.tangent;126            }127128            if (comptime result_mode == .downstream) return self.next.finish(values);129130            var program = try self.next.finish(values);131            errdefer program.deinit();132133            const differentiated_parameters = try differentiatedParameters(&program, self.source, self.options);134            const tangent_parameter_count = differentiated_parameters.len;135136            return .{137                .program = program,138                .differentiated_parameters = differentiated_parameters,139                .primal_parameter_count = self.source.parameters.len,140                .tangent_parameter_count = tangent_parameter_count,141                .primal_output_count = outputs.len,142                .tangent_output_count = outputs.len,143            };144        }145146        pub fn builderHandle(self: *@This()) *trace.Builder {147            return self.next.builderHandle();148        }149150        fn parameterValue(self: *@This(), parameter_info: program_mod.Parameter) !Value {151            if (self.parameter_values == null) try self.prepareParameters();152            return self.parameter_values.?[parameter_info.index];153        }154155        fn prepareParameters(self: *@This()) !void {156            const values = try self.builderHandle().arena.allocator().alloc(Value, self.source.parameters.len);157            for (self.source.parameters) |id| {158                const op = &self.source.operations[id.index];159                const parameter_info = op.kind.parameter;160                values[parameter_info.index].primal = try self.next.bind(op, &.{});161            }162            for (self.source.parameters) |id| {163                const op = &self.source.operations[id.index];164                const parameter_info = op.kind.parameter;165                values[parameter_info.index].tangent = if (isWrt(self.options.wrt, parameter_info.index)) tangent: {166                    if (!op.result.dtype.isFloat()) return error.NonDifferentiableParameter;167                    break :tangent try self.builderHandle().inputTyped(op.result);168                } else try self.emitZeros(op.result);169            }170            self.parameter_values = values;171        }172173        fn compareValue(self: *@This(), op: *const program_mod.Operation, lhs: Value, rhs: Value) !Value {174            return .{175                .primal = try self.next.bind(op, &.{ lhs.primal, rhs.primal }),176                .tangent = try self.emitZeros(op.result),177            };178        }179180        fn selectValue(self: *@This(), op: *const program_mod.Operation, pred: Value, on_true: Value, on_false: Value) !Value {181            return .{182                .primal = try self.next.bind(op, &.{ pred.primal, on_true.primal, on_false.primal }),183                .tangent = try emit.select(self, pred.primal, on_true.tangent, on_false.tangent),184            };185        }186187        fn constantValue(self: *@This(), op: *const program_mod.Operation) !Value {188            return .{189                .primal = try self.next.bind(op, &.{}),190                .tangent = try self.emitZeros(op.result),191            };192        }193194        fn iotaValue(self: *@This(), op: *const program_mod.Operation) !Value {195            return .{196                .primal = try self.next.bind(op, &.{}),197                .tangent = try self.emitZeros(op.result),198            };199        }200201        fn unaryValue(self: *@This(), op: *const program_mod.Operation, unary_info: program_mod.UnaryOp, input: Value) !Value {202            const primal = try self.next.bind(op, &.{input.primal});203            return .{204                .primal = primal,205                .tangent = try unaryJvp(self, unary_info.op, input.primal, primal, input.tangent),206            };207        }208209        fn binaryValue(self: *@This(), op: *const program_mod.Operation, binary_info: program_mod.BinaryOp, lhs: Value, rhs: Value) !Value {210            const primal = try self.next.bind(op, &.{ lhs.primal, rhs.primal });211            return .{212                .primal = primal,213                .tangent = try binaryJvp(self, binary_info.op, lhs.primal, rhs.primal, lhs.tangent, rhs.tangent),214            };215        }216217        fn broadcastValue(self: *@This(), op: *const program_mod.Operation, input: Value) !Value {218            return .{219                .primal = try self.next.bind(op, &.{input.primal}),220                .tangent = try self.emitBroadcast(input.tangent, op.result.dims),221            };222        }223224        fn broadcastInDimValue(self: *@This(), op: *const program_mod.Operation, broadcast_info: program_mod.BroadcastInDim, input: Value) !Value {225            return .{226                .primal = try self.next.bind(op, &.{input.primal}),227                .tangent = try self.emitBroadcastInDim(input.tangent, op.result.dims, broadcast_info.broadcast_dims),228            };229        }230231        fn reshapeValue(self: *@This(), op: *const program_mod.Operation, input: Value) !Value {232            return .{233                .primal = try self.next.bind(op, &.{input.primal}),234                .tangent = try self.emitReshape(input.tangent, op.result.dims),235            };236        }237238        fn transposeValue(self: *@This(), op: *const program_mod.Operation, transpose_info: program_mod.Transpose, input: Value) !Value {239            return .{240                .primal = try self.next.bind(op, &.{input.primal}),241                .tangent = try self.emitTranspose(input.tangent, transpose_info.permutation),242            };243        }244245        fn reduceValue(self: *@This(), op: *const program_mod.Operation, reduce_info: program_mod.Reduce, input: Value, init: Value) !Value {246            const primal = try self.next.bind(op, &.{ input.primal, init.primal });247            const tangent = switch (reduce_info.reducer) {248                .sum => try self.emitReduce(input.tangent, init.tangent, .sum, reduce_info.dimensions),249                .max, .min => blk: {250                    const input_active = !input.tangent.isStructuralZero();251                    const init_active = !init.tangent.isStructuralZero();252                    if (!input_active and !init_active) break :blk try self.emitZeros(primal.ty);253254                    var input_term: ?trace.Value = null;255                    if (input_active) {256                        const allocator = self.builderHandle().arena.allocator();257                        const input_dims = input.primal.ty.dims;258                        const kept = try allocator.alloc(i64, input_dims.len - reduce_info.dimensions.len);259                        var kept_index: usize = 0;260                        for (0..input_dims.len) |axis| {261                            var reduced = false;262                            for (reduce_info.dimensions) |dimension| {263                                if (dimension == @as(i64, @intCast(axis))) reduced = true;264                            }265                            if (reduced) continue;266                            kept[kept_index] = @intCast(axis);267                            kept_index += 1;268                        }269                        const expanded = try self.emitBroadcastInDim(primal, input_dims, kept);270                        const mask = try self.emitCompare(.eq, input.primal, expanded);271                        const zeros = try self.emitZeros(input.primal.ty);272                        const masked = try self.emitSelect(mask, input.tangent, zeros);273                        const zero_init = try self.emitZeros(init.tangent.ty);274                        input_term = try self.emitReduce(masked, zero_init, .sum, reduce_info.dimensions);275                    }276277                    var init_term: ?trace.Value = null;278                    if (init_active) {279                        const init_primal = try self.emitToShape(init.primal, primal.ty);280                        const init_tangent = try self.emitToShape(init.tangent, primal.ty);281                        const init_mask = try self.emitCompare(.eq, primal, init_primal);282                        const zeros = try self.emitZeros(primal.ty);283                        init_term = try self.emitSelect(init_mask, init_tangent, zeros);284                    }285286                    break :blk if (input_term) |input_value|287                        if (init_term) |init_value|288                            try self.emitBinary(.add, input_value, init_value)289                        else290                            input_value291                    else292                        init_term.?;293                },294            };295            return .{296                .primal = primal,297                .tangent = tangent,298            };299        }300301        fn gatherValue(self: *@This(), op: *const program_mod.Operation, gather_info: program_mod.Gather, input: Value, indices: Value) !Value {302            return .{303                .primal = try self.next.bind(op, &.{ input.primal, indices.primal }),304                .tangent = try self.emitGather(input.tangent, indices.primal, gather_info.axis),305            };306        }307308        fn scatterAddValue(309            self: *@This(),310            op: *const program_mod.Operation,311            scatter_add: program_mod.ScatterAdd,312            input: Value,313            indices: Value,314            updates: Value,315        ) !Value {316            return .{317                .primal = try self.next.bind(op, &.{ input.primal, indices.primal, updates.primal }),318                .tangent = try self.emitScatterAdd(input.tangent, indices.primal, updates.tangent, scatter_add.axis),319            };320        }321322        fn sparseCrossEntropyValue(323            self: *@This(),324            op: *const program_mod.Operation,325            sparse_cross_entropy: program_mod.SparseCrossEntropy,326            logits: Value,327            targets: Value,328        ) !Value {329            const primal = try self.next.bind(op, &.{ logits.primal, targets.primal });330331            const logits_ty = logits.primal.ty;332            const axis = sparse_cross_entropy.axis;333            const class_axes = [_]i64{axis};334            const scalar_ty = trace.Type.scalar(logits_ty.dtype);335336            const max_init = try self.emitFloatLowest(logits_ty.dtype);337            const row_max = try emit.reduce(self, logits.primal, max_init, .max, class_axes[0..]);338            const row_max_full = try self.emitExpandAxis(row_max, logits_ty.dims, axis);339            const shifted = try emit.binary(self, .sub, logits.primal, row_max_full);340            const exponentials = try emit.unary(self, .exp, shifted);341            const sum_init = try self.emitZeros(scalar_ty);342            const denominator = try emit.reduce(self, exponentials, sum_init, .sum, class_axes[0..]);343            const denominator_full = try self.emitExpandAxis(denominator, logits_ty.dims, axis);344            const softmax = try emit.binary(self, .div, exponentials, denominator_full);345346            const weighted = try emit.binary(self, .mul, softmax, logits.tangent);347            const weighted_sum = try emit.reduce(self, weighted, sum_init, .sum, class_axes[0..]);348349            const positions = try emit.iota(self, .i32, logits_ty.dims, axis);350            const target_positions = try self.emitExpandAxis(targets.primal, logits_ty.dims, axis);351            const mask = try emit.compare(self, .eq, positions, target_positions);352            const tangent_zeros = try self.emitZeros(logits_ty);353            const masked_tangent = try emit.select(self, mask, logits.tangent, tangent_zeros);354            const target_tangent = try emit.reduce(self, masked_tangent, sum_init, .sum, class_axes[0..]);355356            return .{357                .primal = primal,358                .tangent = try emit.binary(self, .sub, weighted_sum, target_tangent),359            };360        }361362        fn emitExpandAxis(self: *@This(), value: trace.Value, result_dims: []const trace.Dim, axis: i64) !trace.Value {363            const allocator = self.builderHandle().arena.allocator();364            const mapping = try allocator.alloc(i64, value.ty.rank());365            var out: usize = 0;366            for (0..result_dims.len) |position| {367                if (position == @as(usize, @intCast(axis))) continue;368                mapping[out] = @intCast(position);369                out += 1;370            }371            const expanded_dims = try allocator.alloc(trace.Dim, result_dims.len);372            for (result_dims, expanded_dims) |dim, *slot| slot.* = dim;373            return emit.broadcastInDim(self, value, expanded_dims, mapping);374        }375376        fn emitFloatLowest(self: *@This(), dtype: trace.DType) !trace.Value {377            const scalar_ty = trace.Type.scalar(dtype);378            return switch (dtype) {379                .f16 => emit.fullFloat(self, scalar_ty, -std.math.floatMax(f16)),380                .bf16, .f32 => emit.fullFloat(self, scalar_ty, -std.math.floatMax(f32)),381                .f64 => emit.fullFloat(self, scalar_ty, -std.math.floatMax(f64)),382                else => error.NonFloatDType,383            };384        }385386        fn dotGeneralValue(self: *@This(), op: *const program_mod.Operation, dot: program_mod.DotGeneral, lhs: Value, rhs: Value) !Value {387            const primal = try self.next.bind(op, &.{ lhs.primal, rhs.primal });388            const left = try self.emitDotGeneral(389                lhs.tangent,390                rhs.primal,391                dot.lhs_contract,392                dot.rhs_contract,393                dot.lhs_batch,394                dot.rhs_batch,395            );396            const right = try self.emitDotGeneral(397                lhs.primal,398                rhs.tangent,399                dot.lhs_contract,400                dot.rhs_contract,401                dot.lhs_batch,402                dot.rhs_batch,403            );404            return .{405                .primal = primal,406                .tangent = try self.emitBinary(.add, left, right),407            };408        }409410        fn emitZeros(self: *@This(), ty: trace.Type) !trace.Value {411            return emit.zeros(self, ty);412        }413414        fn emitToShape(self: *@This(), value: trace.Value, ty: trace.Type) !trace.Value {415            if (value.ty.dtype != ty.dtype) return error.DTypeMismatch;416            if (types.sameDims(value.ty.dims, ty.dims)) return value;417            if (value.ty.rank() == 0) return self.emitBroadcastInDim(value, ty.dims, &.{});418            return error.ShapeMismatch;419        }420421        fn emitCompare(self: *@This(), direction: program_mod.CompareDirection, lhs: trace.Value, rhs: trace.Value) !trace.Value {422            return emit.compare(self, direction, lhs, rhs);423        }424425        fn emitSelect(self: *@This(), pred: trace.Value, on_true: trace.Value, on_false: trace.Value) !trace.Value {426            return emit.select(self, pred, on_true, on_false);427        }428429        fn emitFullFloat(self: *@This(), ty: trace.Type, fill_value: f64) !trace.Value {430            return emit.fullFloat(self, ty, fill_value);431        }432433        fn emitConstantBytes(self: *@This(), ty: trace.Type, payload: []const u8) !trace.Value {434            return emit.constantBytes(self, ty, payload);435        }436437        fn emitUnary(self: *@This(), op_kind: program_mod.Unary, input: trace.Value) !trace.Value {438            return emit.unary(self, op_kind, input);439        }440441        fn emitBinary(self: *@This(), op_kind: program_mod.Binary, lhs: trace.Value, rhs: trace.Value) !trace.Value {442            return emit.binary(self, op_kind, lhs, rhs);443        }444445        fn emitBroadcast(self: *@This(), input: trace.Value, result_dims: []const trace.Dim) !trace.Value {446            return emit.broadcast(self, input, result_dims);447        }448449        fn emitBroadcastInDim(self: *@This(), input: trace.Value, result_dims: []const trace.Dim, broadcast_dims: []const i64) !trace.Value {450            return emit.broadcastInDim(self, input, result_dims, broadcast_dims);451        }452453        fn emitReshape(self: *@This(), input: trace.Value, new_dims: []const trace.Dim) !trace.Value {454            return emit.reshape(self, input, new_dims);455        }456457        fn emitTranspose(self: *@This(), input: trace.Value, permutation: []const i64) !trace.Value {458            return emit.transpose(self, input, permutation);459        }460461        fn emitReduce(self: *@This(), input: trace.Value, init: trace.Value, reducer: program_mod.Reducer, dimensions: []const i64) !trace.Value {462            return emit.reduce(self, input, init, reducer, dimensions);463        }464465        fn emitGather(self: *@This(), input: trace.Value, indices: trace.Value, axis: i64) !trace.Value {466            return emit.gather(self, input, indices, axis);467        }468469        fn emitScatterAdd(self: *@This(), input: trace.Value, indices: trace.Value, updates: trace.Value, axis: i64) !trace.Value {470            return emit.scatterAdd(self, input, indices, updates, axis);471        }472473        fn emitDotGeneral(474            self: *@This(),475            lhs: trace.Value,476            rhs: trace.Value,477            lhs_contract: []const i64,478            rhs_contract: []const i64,479            lhs_batch: []const i64,480            rhs_batch: []const i64,481        ) !trace.Value {482            return emit.dotGeneral(self, lhs, rhs, lhs_contract, rhs_contract, lhs_batch, rhs_batch);483        }484    };485}486487pub fn linearize(488    allocator: std.mem.Allocator,489    source: *const program_mod.Program,490    options: LinearizeOptions,491) !Linearization {492    if (source.containsScan()) {493        var expanded = try unroll.apply(allocator, source);494        defer expanded.deinit();495        return linearize(allocator, &expanded, options);496    }497498    try validate(source, options);499500    var builder = try trace.Builder.init(allocator, source.name);501    errdefer builder.deinit();502503    const graph = interpret.Graph{ .builder = &builder };504    return linearizeInto(allocator, source, options, graph);505}506507pub fn linearizeWith(508    allocator: std.mem.Allocator,509    source: *const program_mod.Program,510    options: LinearizeOptions,511    hooks: anytype,512) !Linearization {513    if (source.containsScan()) {514        var expanded = try unroll.apply(allocator, source);515        defer expanded.deinit();516        return linearizeWith(allocator, &expanded, options, hooks);517    }518519    try validate(source, options);520521    var builder = try trace.Builder.init(allocator, source.name);522    errdefer builder.deinit();523524    const graph = interpret.Graph{ .builder = &builder };525    const generated = hook.attach("linearize", hooks, graph);526    const linear = semantics(source, generated, options);527    return interpret.run(allocator, source, hook.attach("jvp", hooks, linear));528}529530fn linearizeInto(531    allocator: std.mem.Allocator,532    source: *const program_mod.Program,533    options: LinearizeOptions,534    initial: anytype,535) !Linearization {536    return interpret.run(allocator, source, semantics(source, initial, options));537}538539pub fn jvp(540    allocator: std.mem.Allocator,541    source: *const program_mod.Program,542    options: JvpOptions,543) !program_mod.Program {544    const linearized = try linearize(allocator, source, options);545    return linearized.program;546}547548pub fn jvpWith(549    allocator: std.mem.Allocator,550    source: *const program_mod.Program,551    options: JvpOptions,552    hooks: anytype,553) !program_mod.Program {554    const linearized = try linearizeWith(allocator, source, options, hooks);555    return linearized.program;556}557558fn unaryJvp(layer: anytype, op: program_mod.Unary, input: trace.Value, primal: trace.Value, tangent: trace.Value) !trace.Value {559    if (!input.ty.dtype.isFloat()) return error.UnsupportedDerivative;560    return switch (op) {561        .neg => layer.emitUnary(.neg, tangent),562        .exp => layer.emitBinary(.mul, primal, tangent),563        .log => layer.emitBinary(.div, tangent, input),564        .sqrt => blk: {565            const two = try layer.emitFullFloat(input.ty, 2.0);566            break :blk layer.emitBinary(.div, tangent, try layer.emitBinary(.mul, two, primal));567        },568        .tanh => blk: {569            const one = try layer.emitFullFloat(input.ty, 1.0);570            const square = try layer.emitBinary(.mul, primal, primal);571            break :blk layer.emitBinary(.mul, tangent, try layer.emitBinary(.sub, one, square));572        },573        .sin => layer.emitBinary(.mul, tangent, try layer.emitUnary(.cos, input)),574        .cos => layer.emitUnary(.neg, try layer.emitBinary(.mul, tangent, try layer.emitUnary(.sin, input))),575        .abs => blk: {576            const zeros = try layer.emitZeros(input.ty);577            const non_negative = try layer.emitCompare(.ge, input, zeros);578            break :blk layer.emitSelect(non_negative, tangent, try layer.emitUnary(.neg, tangent));579        },580        .tan => blk: {581            const cosine = try layer.emitUnary(.cos, input);582            break :blk layer.emitBinary(.div, tangent, try layer.emitBinary(.mul, cosine, cosine));583        },584    };585}586587fn binaryJvp(588    layer: anytype,589    op: program_mod.Binary,590    lhs: trace.Value,591    rhs: trace.Value,592    lhs_tangent: trace.Value,593    rhs_tangent: trace.Value,594) !trace.Value {595    if (!lhs.ty.dtype.isFloat()) return error.UnsupportedDerivative;596    return switch (op) {597        .add => layer.emitBinary(.add, lhs_tangent, rhs_tangent),598        .sub => layer.emitBinary(.sub, lhs_tangent, rhs_tangent),599        .mul => layer.emitBinary(.add, try layer.emitBinary(.mul, lhs_tangent, rhs), try layer.emitBinary(.mul, lhs, rhs_tangent)),600        .div => blk: {601            const numerator = try layer.emitBinary(.sub, try layer.emitBinary(.mul, lhs_tangent, rhs), try layer.emitBinary(.mul, lhs, rhs_tangent));602            const denominator = try layer.emitBinary(.mul, rhs, rhs);603            break :blk layer.emitBinary(.div, numerator, denominator);604        },605        .max => blk: {606            const mask = try layer.emitCompare(.ge, lhs, rhs);607            break :blk layer.emitSelect(mask, lhs_tangent, rhs_tangent);608        },609        .min => blk: {610            const mask = try layer.emitCompare(.le, lhs, rhs);611            break :blk layer.emitSelect(mask, lhs_tangent, rhs_tangent);612        },613        .pow => blk: {614            const lhs_active = !lhs_tangent.isStructuralZero();615            const rhs_active = !rhs_tangent.isStructuralZero();616            if (!lhs_active and !rhs_active) break :blk layer.emitZeros(lhs.ty);617618            var base_term: ?trace.Value = null;619            if (lhs_active) {620                const one = try layer.emitFullFloat(lhs.ty, 1.0);621                const exponent_minus_one = try layer.emitBinary(.sub, rhs, one);622                const power_step_down = try layer.emitBinary(.pow, lhs, exponent_minus_one);623                base_term = try layer.emitBinary(624                    .mul,625                    try layer.emitBinary(.mul, rhs, power_step_down),626                    lhs_tangent,627                );628            }629630            var exponent_term: ?trace.Value = null;631            if (rhs_active) {632                exponent_term = try layer.emitBinary(633                    .mul,634                    try layer.emitBinary(.mul, try layer.emitBinary(.pow, lhs, rhs), try layer.emitUnary(.log, lhs)),635                    rhs_tangent,636                );637            }638639            break :blk if (base_term) |base|640                if (exponent_term) |exponent|641                    try layer.emitBinary(.add, base, exponent)642                else643                    base644            else645                exponent_term.?;646        },647    };648}649650fn differentiatedParameters(651    program: *program_mod.Program,652    source: *const program_mod.Program,653    options: LinearizeOptions,654) ![]const usize {655    const result = try program.arena.allocator().alloc(usize, differentiatedParameterCount(source, options.wrt));656    var out: usize = 0;657    for (source.parameters) |id| {658        const parameter = source.operation(id).kind.parameter;659        if (isWrt(options.wrt, parameter.index)) {660            result[out] = parameter.index;661            out += 1;662        }663    }664    return result;665}666667fn differentiatedParameterCount(source: *const program_mod.Program, wrt: []const usize) usize {668    var count: usize = 0;669    for (source.parameters) |id| {670        const parameter = source.operation(id).kind.parameter;671        if (isWrt(wrt, parameter.index)) count += 1;672    }673    return count;674}675676pub fn validate(source: *const program_mod.Program, options: LinearizeOptions) !void {677    for (options.wrt, 0..) |index, offset| {678        if (index >= source.parameters.len) return error.ParameterOutOfRange;679        for (options.wrt[0..offset]) |seen| {680            if (seen == index) return error.DuplicateParameter;681        }682    }683}684685fn isWrt(wrt: []const usize, index: usize) bool {686    for (wrt) |item| {687        if (item == index) return true;688    }689    return false;690}691692fn jvpBody(_: *trace.Builder, args: []const trace.Value) !trace.Value {693    return try (try args[0].mul(args[1])).tanh();694}695696fn addZeroBody(builder: *trace.Builder, args: []const trace.Value) !trace.Value {697    const zero = try builder.full(.f32, .{ .lane = 4 }, 0.0);698    return try args[0].add(zero);699}700701fn powInactiveExponentBody(builder: *trace.Builder, args: []const trace.Value) !trace.Value {702    const three = try builder.scalar(.f32, 3.0);703    const exponent = three;704    return try args[0].pow(exponent);705}706707fn reduceMaxWithInitBody(_: *trace.Builder, args: []const trace.Value) !trace.Value {708    return try args[0].reduce(args[1], .max, .lane);709}710711fn reduceMinWithInitBody(_: *trace.Builder, args: []const trace.Value) !trace.Value {712    return try args[0].reduce(args[1], .min, .lane);713}714715fn expectExtremaInitTangentMasked(source: *const program_mod.Program) !void {716    var differentiated = try linearize(std.testing.allocator, source, .{ .wrt = &.{ 0, 1 } });717    defer differentiated.deinit();718719    try std.testing.expectEqual(@as(usize, 4), differentiated.program.parameters.len);720    const init_tangent = differentiated.program.parameters[3];721    var reduce_seeded_by_init_tangent = false;722    var masked_init_tangent = false;723    for (differentiated.program.operations) |op| {724        switch (op.kind) {725            .reduce => |reduce| {726                if (reduce.reducer == .sum and reduce.init.index == init_tangent.index) {727                    reduce_seeded_by_init_tangent = true;728                }729            },730            .select => |select| {731                if (select.on_true.index == init_tangent.index) masked_init_tangent = true;732            },733            else => {},734        }735    }736737    try std.testing.expect(!reduce_seeded_by_init_tangent);738    try std.testing.expect(masked_init_tangent);739}740741const CountDualBinary = struct {742    count: *usize,743744    pub fn bind(self: *@This(), ctx: anytype) !Dual {745        switch (ctx.op.kind) {746            .binary => self.count.* += 1,747            else => {},748        }749        return ctx.default();750    }751};752753const TraceBinaryCounts = struct {754    add: usize = 0,755    mul: usize = 0,756};757758const TraceBinaryCounter = struct {759    counts: *TraceBinaryCounts,760761    pub fn add(self: *@This(), ctx: anytype) !trace.Value {762        self.counts.add += 1;763        return ctx.default();764    }765766    pub fn mul(self: *@This(), ctx: anytype) !trace.Value {767        self.counts.mul += 1;768        return ctx.default();769    }770};771772test "tensor linearize semantics accepts user layers over dual values" {773    var source = try trace.define(std.testing.allocator, "jvp_layer", &.{774        types.spec(.f32, .{ .lane = 4 }),775    }, addZeroBody);776    defer source.deinit();777778    var builder = try trace.Builder.init(std.testing.allocator, source.name);779    errdefer builder.deinit();780781    var count: usize = 0;782    const graph = interpret.Graph{ .builder = &builder };783    const linear = semantics(&source, graph, .{ .wrt = &.{0} });784    var differentiated = try interpret.run(std.testing.allocator, &source, interpret.layer(Dual, linear, CountDualBinary{ .count = &count }));785    defer differentiated.deinit();786787    try std.testing.expectEqual(@as(usize, 1), count);788    try std.testing.expectEqual(@as(usize, 2), differentiated.program.parameters.len);789    try std.testing.expectEqual(@as(usize, 2), differentiated.program.outputs.len);790}791792test "tensor linearize binds generated tangent arithmetic through downstream semantics" {793    var source = try trace.define(std.testing.allocator, "jvp_generated_arithmetic", &.{794        types.spec(.f32, .{ .lane = 4 }),795        types.spec(.f32, .{ .lane = 4 }),796    }, struct {797        fn body(_: *trace.Builder, args: []const trace.Value) !trace.Value {798            return try args[0].mul(args[1]);799        }800    }.body);801    defer source.deinit();802803    var builder = try trace.Builder.init(std.testing.allocator, source.name);804    errdefer builder.deinit();805806    var counts: TraceBinaryCounts = .{};807    const graph = interpret.Graph{ .builder = &builder };808    const counted = interpret.bind(TraceBinaryCounter{ .counts = &counts }).attach(graph);809    var differentiated = try interpret.run(std.testing.allocator, &source, semantics(&source, counted, .{ .wrt = &.{ 0, 1 } }));810    defer differentiated.deinit();811812    try std.testing.expectEqual(@as(usize, 1), counts.add);813    try std.testing.expectEqual(@as(usize, 3), counts.mul);814    try std.testing.expectEqual(@as(usize, 4), differentiated.program.parameters.len);815    try std.testing.expectEqual(@as(usize, 2), differentiated.program.outputs.len);816}817818const GeneratedAddMetadata = struct {819    seen: *usize,820821    pub fn add(self: *@This(), ctx: *transform.Context) !?trace.Value {822        self.seen.* += 1;823        try std.testing.expect(!ctx.isZero(0));824        try std.testing.expect(!ctx.isZero(1));825        try std.testing.expect(ctx.constantPayload(0) == null);826        try std.testing.expect(ctx.constantPayload(1) == null);827        return null;828    }829};830831test "tensor linearize generated arithmetic is safe for downstream rewrite metadata queries" {832    var source = try trace.define(std.testing.allocator, "jvp_generated_metadata", &.{833        types.spec(.f32, .{ .lane = 4 }),834        types.spec(.f32, .{ .lane = 4 }),835    }, struct {836        fn body(_: *trace.Builder, args: []const trace.Value) !trace.Value {837            return try args[0].mul(args[1]);838        }839    }.body);840    defer source.deinit();841842    var builder = try trace.Builder.init(std.testing.allocator, source.name);843    errdefer builder.deinit();844845    var seen: usize = 0;846    const graph = interpret.Graph{ .builder = &builder };847    const rewrite = transform.semantics(&source, graph, GeneratedAddMetadata{ .seen = &seen });848    var differentiated = try interpret.run(std.testing.allocator, &source, semantics(&source, rewrite, .{ .wrt = &.{ 0, 1 } }));849    defer differentiated.deinit();850851    try std.testing.expectEqual(@as(usize, 1), seen);852    try std.testing.expectEqual(@as(usize, 4), differentiated.program.parameters.len);853    try std.testing.expectEqual(@as(usize, 2), differentiated.program.outputs.len);854}855856const DropAddZero = struct {857    pub fn add(_: *@This(), ctx: *transform.Context) !?trace.Value {858        if (ctx.isZero(1)) return ctx.arg(0);859        if (ctx.isZero(0)) return ctx.arg(1);860        return null;861    }862};863864test "tensor linearize semantics delegates primal binds through rewrite layers" {865    var source = try trace.define(std.testing.allocator, "jvp_rewrite_stack", &.{866        types.spec(.f32, .{ .lane = 4 }),867    }, addZeroBody);868    defer source.deinit();869870    var plain = try linearize(std.testing.allocator, &source, .{ .wrt = &.{0} });871    defer plain.deinit();872873    var builder = try trace.Builder.init(std.testing.allocator, source.name);874    errdefer builder.deinit();875876    const graph = interpret.Graph{ .builder = &builder };877    const rewrite = transform.semantics(&source, graph, DropAddZero{});878    var rewritten = try interpret.run(std.testing.allocator, &source, semantics(&source, rewrite, .{ .wrt = &.{0} }));879    defer rewritten.deinit();880881    try std.testing.expect(rewritten.program.operationCount() < plain.program.operationCount());882    try std.testing.expectEqual(@as(u32, 0), rewritten.primalOutputs()[0].index);883    try types.expectExtents(&.{4}, rewritten.program.typeOf(rewritten.tangentOutputs()[0]));884}885886test "tensor linearize exposes primal and tangent structure" {887    var source = try trace.define(std.testing.allocator, "jvp", &.{888        types.spec(.f32, .{ .lane = 4 }),889        types.spec(.f32, .{ .lane = 4 }),890    }, jvpBody);891    defer source.deinit();892893    var differentiated = try linearize(std.testing.allocator, &source, .{ .wrt = &.{ 0, 1 } });894    defer differentiated.deinit();895896    try std.testing.expectEqual(@as(usize, 2), differentiated.primal_parameter_count);897    try std.testing.expectEqual(@as(usize, 2), differentiated.tangent_parameter_count);898    try std.testing.expectEqual(@as(usize, 1), differentiated.primal_output_count);899    try std.testing.expectEqual(@as(usize, 1), differentiated.tangent_output_count);900    try std.testing.expectEqualSlices(usize, &.{ 0, 1 }, differentiated.differentiated_parameters);901    try std.testing.expectEqual(@as(usize, 4), differentiated.program.parameters.len);902    try std.testing.expectEqual(@as(usize, 2), differentiated.program.outputs.len);903    try types.expectExtents(&.{4}, differentiated.program.typeOf(differentiated.primalOutputs()[0]));904    try types.expectExtents(&.{4}, differentiated.program.typeOf(differentiated.tangentOutputs()[0]));905}906907test "tensor linearize omits inactive pow exponent branch" {908    var source = try trace.define(std.testing.allocator, "jvp_pow_inactive_exponent", &.{909        types.spec(.f32, .{ .lane = 4 }),910    }, powInactiveExponentBody);911    defer source.deinit();912913    var differentiated = try linearize(std.testing.allocator, &source, .{ .wrt = &.{0} });914    defer differentiated.deinit();915916    var logs: usize = 0;917    for (differentiated.program.operations) |op| {918        switch (op.kind) {919            .unary => |unary_op| {920                if (unary_op.op == .log) logs += 1;921            },922            else => {},923        }924    }925926    try std.testing.expectEqual(@as(usize, 0), logs);927    try types.expectExtents(&.{4}, differentiated.program.typeOf(differentiated.tangentOutputs()[0]));928}929930test "tensor linearize masks reduce max init tangents" {931    var source = try trace.define(std.testing.allocator, "jvp_reduce_max_init", &.{932        types.spec(.f32, .{ .lane = 4 }),933        types.spec(.f32, .{}),934    }, reduceMaxWithInitBody);935    defer source.deinit();936937    try expectExtremaInitTangentMasked(&source);938}939940test "tensor linearize masks reduce min init tangents" {941    var source = try trace.define(std.testing.allocator, "jvp_reduce_min_init", &.{942        types.spec(.f32, .{ .lane = 4 }),943        types.spec(.f32, .{}),944    }, reduceMinWithInitBody);945    defer source.deinit();946947    try expectExtremaInitTangentMasked(&source);948}949950test "tensor jvp is a linearize convenience wrapper" {951    var source = try trace.define(std.testing.allocator, "jvp_wrapper", &.{952        types.spec(.f32, .{ .lane = 4 }),953        types.spec(.f32, .{ .lane = 4 }),954    }, jvpBody);955    defer source.deinit();956957    var differentiated = try jvp(std.testing.allocator, &source, .{ .wrt = &.{ 0, 1 } });958    defer differentiated.deinit();959960    try std.testing.expectEqual(@as(usize, 4), differentiated.parameters.len);961    try std.testing.expectEqual(@as(usize, 2), differentiated.outputs.len);962    try types.expectExtents(&.{4}, differentiated.typeOf(differentiated.outputs[0]));963    try types.expectExtents(&.{4}, differentiated.typeOf(differentiated.outputs[1]));964}965966test "tensor jvpWith routes linearize generated ops through user semantics" {967    var source = try trace.define(std.testing.allocator, "jvp_with_generated_arithmetic", &.{968        types.spec(.f32, .{ .lane = 4 }),969        types.spec(.f32, .{ .lane = 4 }),970    }, struct {971        fn body(_: *trace.Builder, args: []const trace.Value) !trace.Value {972            return try args[0].mul(args[1]);973        }974    }.body);975    defer source.deinit();976977    var counts: TraceBinaryCounts = .{};978    var differentiated = try jvpWith(979        std.testing.allocator,980        &source,981        .{ .wrt = &.{ 0, 1 } },982        .{ .linearize = interpret.bind(TraceBinaryCounter{ .counts = &counts }) },983    );984    defer differentiated.deinit();985986    try std.testing.expectEqual(@as(usize, 1), counts.add);987    try std.testing.expectEqual(@as(usize, 3), counts.mul);988    try std.testing.expectEqual(@as(usize, 4), differentiated.parameters.len);989    try std.testing.expectEqual(@as(usize, 2), differentiated.outputs.len);990    try types.expectExtents(&.{4}, differentiated.typeOf(differentiated.outputs[0]));991    try types.expectExtents(&.{4}, differentiated.typeOf(differentiated.outputs[1]));992}993994fn doublingCallBody(builder: *trace.Builder, args: []const trace.Value) !trace.Value {995    _ = builder;996    return args[0].builder.customCall("accy.custom.double", 1, &.{args[0]}, args[0].ty);997}998999test "tensor linearize rejects opaque custom calls naming the contract" {1000    var source = try trace.define(std.testing.allocator, "jvp_opaque_custom", &.{1001        types.spec(.f32, .{ .lane = 4 }),1002    }, doublingCallBody);1003    defer source.deinit();10041005    try std.testing.expectError(1006        error.CustomCallRequiresJvpContract,1007        linearize(std.testing.allocator, &source, .{ .wrt = &.{0} }),1008    );1009}10101011const DoubleJvpRule = struct {1012    applied: *usize,10131014    pub fn bind(self: *@This(), ctx: anytype) !Dual {1015        switch (ctx.op.kind) {1016            .custom_call => |custom| {1017                if (std.mem.eql(u8, custom.target, "accy.custom.double")) {1018                    self.applied.* += 1;1019                    const builder = ctx.builderHandle();1020                    return .{1021                        .primal = try builder.customCall(custom.target, custom.version, &.{ctx.args[0].primal}, ctx.op.result),1022                        .tangent = try builder.customCall(custom.target, custom.version, &.{ctx.args[0].tangent}, ctx.op.result),1023                    };1024                }1025            },1026            else => {},1027        }1028        return ctx.default();1029    }1030};10311032test "tensor linearize accepts custom call jvp contracts through user layers" {1033    var source = try trace.define(std.testing.allocator, "jvp_custom_contract", &.{1034        types.spec(.f32, .{ .lane = 4 }),1035    }, doublingCallBody);1036    defer source.deinit();10371038    var builder = try trace.Builder.init(std.testing.allocator, source.name);1039    errdefer builder.deinit();10401041    var applied: usize = 0;1042    const graph = interpret.Graph{ .builder = &builder };1043    const linear = semantics(&source, graph, .{ .wrt = &.{0} });1044    var differentiated = try interpret.run(1045        std.testing.allocator,1046        &source,1047        interpret.layer(Dual, linear, DoubleJvpRule{ .applied = &applied }),1048    );1049    defer differentiated.deinit();10501051    try std.testing.expectEqual(@as(usize, 1), applied);1052    try std.testing.expectEqual(@as(usize, 2), differentiated.program.parameters.len);1053    try std.testing.expectEqual(@as(usize, 2), differentiated.program.outputs.len);10541055    var custom_calls: usize = 0;1056    for (differentiated.program.operations) |op| {1057        switch (op.kind) {1058            .custom_call => custom_calls += 1,1059            else => {},1060        }1061    }1062    try std.testing.expectEqual(@as(usize, 2), custom_calls);1063}10641065test "tensor jvpWith accepts custom call jvp contracts through public hooks" {1066    var source = try trace.define(std.testing.allocator, "jvp_public_custom_contract", &.{1067        types.spec(.f32, .{ .lane = 4 }),1068    }, doublingCallBody);1069    defer source.deinit();10701071    var applied: usize = 0;1072    var differentiated = try jvpWith(1073        std.testing.allocator,1074        &source,1075        .{ .wrt = &.{0} },1076        .{ .jvp = interpret.bind(DoubleJvpRule{ .applied = &applied }) },1077    );1078    defer differentiated.deinit();10791080    try std.testing.expectEqual(@as(usize, 1), applied);1081    try std.testing.expectEqual(@as(usize, 2), differentiated.parameters.len);1082    try std.testing.expectEqual(@as(usize, 2), differentiated.outputs.len);10831084    var custom_calls: usize = 0;1085    for (differentiated.operations) |op| {1086        switch (op.kind) {1087            .custom_call => custom_calls += 1,1088            else => {},1089        }1090    }1091    try std.testing.expectEqual(@as(usize, 2), custom_calls);1092}

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