lib/accy/src/tensor/reverse.zig
daab053ee43316e1809a84551d573ddd1e5bf3d2
1 const std = @import("std");
2 const autodiff = @import("autodiff.zig");
3 const emit = @import("emit.zig");
4 const interpret = @import("interpret/root.zig");
5 const program_mod = @import("program.zig");
6 const trace = @import("trace/root.zig");
7 const transform = @import("transform.zig");
8 const types = @import("type/root.zig");
9
10 pub const Options = struct {
11 keep_primal_outputs: bool = false,
12 };
13
14 pub const Pullback = struct {
15 program: program_mod.Program,
16 differentiated_parameters: []const usize,
17 primal_parameter_count: usize,
18 seed_parameter_count: usize,
19 primal_output_count: usize,
20 cotangent_output_count: usize,
21
22 pub fn deinit(self: *Pullback) void {
23 self.program.deinit();
24 }
25 };
26
27 pub fn pullback(allocator: std.mem.Allocator, linearized: *const autodiff.Linearization, options: Options) !Pullback {
28 var builder = try trace.Builder.init(allocator, linearized.program.name);
29 errdefer builder.deinit();
30
31 const graph = interpret.Graph{ .builder = &builder };
32 return pullbackWith(allocator, linearized, options, graph);
33 }
34
35 pub const NoVjpRules = struct {};
36
37 fn RuleHandle(comptime Storage: type) type {
38 const Child = @typeInfo(Storage).pointer.child;
39 return switch (@typeInfo(Child)) {
40 .pointer => Child,
41 else => Storage,
42 };
43 }
44
45 fn ruleHandle(storage: anytype) RuleHandle(@TypeOf(storage)) {
46 const Child = @typeInfo(@TypeOf(storage)).pointer.child;
47 return switch (@typeInfo(Child)) {
48 .pointer => storage.*,
49 else => storage,
50 };
51 }
52
53 pub fn pullbackWith(allocator: std.mem.Allocator, linearized: *const autodiff.Linearization, options: Options, initial: anytype) !Pullback {
54 return pullbackWithRules(allocator, linearized, options, initial, NoVjpRules{});
55 }
56
57 pub fn pullbackWithRules(
58 allocator: std.mem.Allocator,
59 linearized: *const autodiff.Linearization,
60 options: Options,
61 initial: anytype,
62 rules: anytype,
63 ) !Pullback {
64 var layer = TransposeLayer(@TypeOf(initial)){ .next = initial };
65
66 const active = try allocator.alloc(bool, linearized.program.valueCount());
67 defer allocator.free(active);
68 try markActive(linearized, active);
69
70 const needed = try allocator.alloc(bool, linearized.program.valueCount());
71 defer allocator.free(needed);
72 for (needed) |*slot| slot.* = false;
73 try markResiduals(&linearized.program, active, needed);
74 for (linearized.program.parameters) |id| {
75 const op = linearized.program.operation(id);
76 if (op.kind.parameter.index < linearized.primal_parameter_count) needed[id.index] = true;
77 }
78 if (options.keep_primal_outputs) {
79 for (linearized.primalOutputs()) |id| {
80 try markResidual(&linearized.program, active, needed, id);
81 }
82 }
83
84 const residuals = try allocator.alloc(?trace.Value, linearized.program.valueCount());
85 defer allocator.free(residuals);
86 for (residuals) |*slot| slot.* = null;
87 try replayResiduals(allocator, &layer, &linearized.program, active, needed, residuals);
88
89 const cotangents = try allocator.alloc(?trace.Value, linearized.program.valueCount());
90 defer allocator.free(cotangents);
91 for (cotangents) |*slot| slot.* = null;
92
93 for (linearized.tangentOutputs()) |id| {
94 const seed = try layer.builderHandle().inputTyped(linearized.program.typeOf(id));
95 try addCotangent(&layer, cotangents, id, seed);
96 }
97
98 var mutable_rules = rules;
99 try transposeActive(&layer, &linearized.program, active, residuals, cotangents, ruleHandle(&mutable_rules));
100
101 const primal_output_count: usize = if (options.keep_primal_outputs) linearized.primal_output_count else 0;
102 const outputs = try allocator.alloc(trace.Value, primal_output_count + linearized.tangent_parameter_count);
103 defer allocator.free(outputs);
104 var output_index: usize = 0;
105 if (options.keep_primal_outputs) {
106 for (linearized.primalOutputs()) |id| {
107 outputs[output_index] = requireResidual(residuals, id);
108 output_index += 1;
109 }
110 }
111 for (linearized.program.parameters) |id| {
112 const op = linearized.program.operation(id);
113 const parameter = op.kind.parameter;
114 if (parameter.index < linearized.primal_parameter_count) continue;
115 outputs[output_index] = cotangents[id.index] orelse try emit.zeros(&layer, op.result);
116 output_index += 1;
117 }
118
119 const differentiated_parameters = try layer.builderHandle().arena.allocator().dupe(usize, linearized.differentiated_parameters);
120
121 return .{
122 .program = try layer.next.finish(outputs),
123 .differentiated_parameters = differentiated_parameters,
124 .primal_parameter_count = linearized.primal_parameter_count,
125 .seed_parameter_count = linearized.tangent_output_count,
126 .primal_output_count = primal_output_count,
127 .cotangent_output_count = linearized.tangent_parameter_count,
128 };
129 }
130
131 fn TransposeLayer(comptime Next: type) type {
132 return struct {
133 next: Next,
134
135 pub fn builderHandle(self: *@This()) *trace.Builder {
136 return self.next.builderHandle();
137 }
138 };
139 }
140
141 fn markActive(linearized: *const autodiff.Linearization, active: []bool) !void {
142 for (linearized.program.operations) |op| {
143 active[op.id.index] = switch (op.kind) {
144 .parameter => |parameter| parameter.index >= linearized.primal_parameter_count,
145 .constant, .iota => false,
146 .unary => |unary| active[unary.input.index],
147 .binary => |binary| active[binary.lhs.index] or active[binary.rhs.index],
148 .broadcast => |broadcast| active[broadcast.input.index],
149 .broadcast_in_dim => |broadcast| active[broadcast.input.index],
150 .reshape => |reshape| active[reshape.input.index],
151 .transpose => |transpose| active[transpose.input.index],
152 .compare => false,
153 .select => |select| active[select.on_true.index] or active[select.on_false.index],
154 .custom_call => |custom| blk: {
155 var any = false;
156 for (custom.operands) |operand| {
157 if (active[operand.index]) any = true;
158 }
159 break :blk any;
160 },
161 .reduce => |reduce| active[reduce.input.index] or active[reduce.init.index],
162 .gather => |gather| active[gather.input.index],
163 .scatter_add => |scatter_add| active[scatter_add.input.index] or active[scatter_add.updates.index],
164 .sparse_cross_entropy => |sparse_cross_entropy| active[sparse_cross_entropy.logits.index],
165 .dot_general => |dot| active[dot.lhs.index] or active[dot.rhs.index],
166 .scan => |scan| blk: {
167 var any = false;
168 for (scan.inits) |init_id| {
169 if (active[init_id.index]) any = true;
170 }
171 break :blk any;
172 },
173 .projection => |projection| active[projection.source.index],
174 };
175 }
176 }
177
178 fn replayResiduals(
179 allocator: std.mem.Allocator,
180 layer: anytype,
181 program: *const program_mod.Program,
182 active: []const bool,
183 needed: []const bool,
184 residuals: []?trace.Value,
185 ) !void {
186 try interpret.run(allocator, program, ResidualReplay(@TypeOf(layer)){
187 .layer = layer,
188 .active = active,
189 .needed = needed,
190 .residuals = residuals,
191 });
192 }
193
194 fn ResidualReplay(comptime Layer: type) type {
195 return struct {
196 layer: Layer,
197 active: []const bool,
198 needed: []const bool,
199 residuals: []?trace.Value,
200
201 pub const Value = ?trace.Value;
202 pub const Result = void;
203
204 pub fn operation(self: *@This(), step: *interpret.Step(Value)) !Value {
205 var buffer: [program_mod.max_operation_operands]Value = undefined;
206 return self.bind(step.op, interpret.arguments(Value, step.op, step.values, &buffer));
207 }
208
209 pub fn bind(self: *@This(), op: *const program_mod.Operation, args: []const Value) !Value {
210 if (self.active[op.id.index] or !self.needed[op.id.index]) return null;
211
212 var buffer: [program_mod.max_operation_operands]trace.Value = undefined;
213 const value = try self.layer.next.bind(op, requireArgs(args, &buffer));
214 self.residuals[op.id.index] = value;
215 return value;
216 }
217
218 pub fn finish(_: *@This(), outputs: []const Value) !Result {
219 _ = outputs;
220 }
221
222 pub fn builderHandle(self: *@This()) *trace.Builder {
223 return self.layer.builderHandle();
224 }
225 };
226 }
227
228 fn requireArgs(args: []const ?trace.Value, buffer: *[program_mod.max_operation_operands]trace.Value) []const trace.Value {
229 for (args, 0..) |arg, index| {
230 buffer[index] = arg orelse unreachable;
231 }
232 return buffer[0..args.len];
233 }
234
235 fn markResiduals(program: *const program_mod.Program, active: []const bool, needed: []bool) !void {
236 for (program.operations) |op| {
237 if (!active[op.id.index]) continue;
238 switch (op.kind) {
239 .parameter, .constant, .iota => {},
240 .unary => {},
241 .binary => |binary| try markBinaryResiduals(program, active, needed, binary),
242 .dot_general => |dot| try markDotResiduals(program, active, needed, dot),
243 .compare => {},
244 .select => |select| try markResidual(program, active, needed, select.pred),
245 .gather => |gather| {
246 if (active[gather.input.index]) try markResidual(program, active, needed, gather.indices);
247 },
248 .scatter_add => |scatter_add| {
249 if (active[scatter_add.updates.index]) try markResidual(program, active, needed, scatter_add.indices);
250 },
251 .sparse_cross_entropy => return error.SparseCrossEntropyTransposeUnsupported,
252 .custom_call => |custom| try markCustomCallResiduals(program, active, needed, custom),
253 .broadcast, .broadcast_in_dim, .reshape, .transpose, .reduce => {},
254 .scan, .projection => return error.ScanTransposeUnsupported,
255 }
256 }
257 }
258
259 fn markBinaryResiduals(
260 program: *const program_mod.Program,
261 active: []const bool,
262 needed: []bool,
263 binary: program_mod.BinaryOp,
264 ) !void {
265 const lhs_active = active[binary.lhs.index];
266 const rhs_active = active[binary.rhs.index];
267 switch (binary.op) {
268 .mul => {
269 if (lhs_active and !rhs_active) try markResidual(program, active, needed, binary.rhs);
270 if (rhs_active and !lhs_active) try markResidual(program, active, needed, binary.lhs);
271 },
272 .div => {
273 if (lhs_active and !rhs_active) try markResidual(program, active, needed, binary.rhs);
274 },
275 .add, .sub, .max, .min, .pow => {},
276 }
277 }
278
279 fn markDotResiduals(
280 program: *const program_mod.Program,
281 active: []const bool,
282 needed: []bool,
283 dot: program_mod.DotGeneral,
284 ) !void {
285 const lhs_active = active[dot.lhs.index];
286 const rhs_active = active[dot.rhs.index];
287 if (lhs_active and !rhs_active) try markResidual(program, active, needed, dot.rhs);
288 if (rhs_active and !lhs_active) try markResidual(program, active, needed, dot.lhs);
289 }
290
291 fn markCustomCallResiduals(
292 program: *const program_mod.Program,
293 active: []const bool,
294 needed: []bool,
295 custom: program_mod.CustomCall,
296 ) !void {
297 for (custom.operands) |operand| {
298 if (!active[operand.index]) try markResidual(program, active, needed, operand);
299 }
300 }
301
302 fn markResidual(
303 program: *const program_mod.Program,
304 active: []const bool,
305 needed: []bool,
306 id: program_mod.Id,
307 ) !void {
308 if (active[id.index]) return error.NonlinearTranspose;
309 if (needed[id.index]) return;
310 needed[id.index] = true;
311 const op = program.operation(id);
312 switch (op.kind) {
313 .parameter, .constant, .iota => {},
314 .unary => |unary| try markResidual(program, active, needed, unary.input),
315 .binary => |binary| {
316 try markResidual(program, active, needed, binary.lhs);
317 try markResidual(program, active, needed, binary.rhs);
318 },
319 .broadcast => |broadcast| try markResidual(program, active, needed, broadcast.input),
320 .broadcast_in_dim => |broadcast| try markResidual(program, active, needed, broadcast.input),
321 .reshape => |reshape| try markResidual(program, active, needed, reshape.input),
322 .transpose => |transpose| try markResidual(program, active, needed, transpose.input),
323 .compare => |compare| {
324 try markResidual(program, active, needed, compare.lhs);
325 try markResidual(program, active, needed, compare.rhs);
326 },
327 .select => |select| {
328 try markResidual(program, active, needed, select.pred);
329 try markResidual(program, active, needed, select.on_true);
330 try markResidual(program, active, needed, select.on_false);
331 },
332 .custom_call => |custom| {
333 for (custom.operands) |operand| try markResidual(program, active, needed, operand);
334 },
335 .reduce => |reduce| {
336 try markResidual(program, active, needed, reduce.input);
337 try markResidual(program, active, needed, reduce.init);
338 },
339 .gather => |gather| {
340 try markResidual(program, active, needed, gather.input);
341 try markResidual(program, active, needed, gather.indices);
342 },
343 .scatter_add => |scatter_add| {
344 try markResidual(program, active, needed, scatter_add.input);
345 try markResidual(program, active, needed, scatter_add.indices);
346 try markResidual(program, active, needed, scatter_add.updates);
347 },
348 .sparse_cross_entropy => |sparse_cross_entropy| {
349 try markResidual(program, active, needed, sparse_cross_entropy.logits);
350 try markResidual(program, active, needed, sparse_cross_entropy.targets);
351 },
352 .dot_general => |dot| {
353 try markResidual(program, active, needed, dot.lhs);
354 try markResidual(program, active, needed, dot.rhs);
355 },
356 .scan => |scan| {
357 for (scan.inits) |init_id| try markResidual(program, active, needed, init_id);
358 },
359 .projection => |projection| try markResidual(program, active, needed, projection.source),
360 }
361 }
362
363 fn transposeActive(
364 layer: anytype,
365 program: *const program_mod.Program,
366 active: []const bool,
367 residuals: []const ?trace.Value,
368 cotangents: []?trace.Value,
369 rules: anytype,
370 ) !void {
371 var index = program.operations.len;
372 while (index > 0) {
373 index -= 1;
374 const op = &program.operations[index];
375 const cotangent = cotangents[op.id.index] orelse continue;
376 if (!active[op.id.index]) continue;
377
378 switch (op.kind) {
379 .parameter, .constant, .iota => {},
380 .unary => |unary| try transposeUnary(layer, unary, cotangent, cotangents),
381 .binary => |binary| try transposeBinary(layer, binary, cotangent, active, residuals, cotangents),
382 .broadcast => |broadcast| try transposeBroadcast(layer, program, op, broadcast, cotangent, active, cotangents),
383 .broadcast_in_dim => |broadcast| try transposeBroadcastInDim(layer, program, op, broadcast, cotangent, active, cotangents),
384 .reshape => |reshape| try transposeReshape(layer, program, reshape, cotangent, active, cotangents),
385 .transpose => |transpose| try transposeTranspose(layer, program, transpose, cotangent, active, cotangents),
386 .reduce => |reduce| try transposeReduce(layer, program, reduce, cotangent, active, cotangents),
387 .gather => |gather| try transposeGather(layer, program, gather, cotangent, active, residuals, cotangents),
388 .scatter_add => |scatter_add| try transposeScatterAdd(layer, scatter_add, cotangent, active, residuals, cotangents),
389 .sparse_cross_entropy => return error.SparseCrossEntropyTransposeUnsupported,
390 .dot_general => |dot| try transposeDot(layer, program, dot, cotangent, active, residuals, cotangents),
391 .compare => {},
392 .select => |select| try transposeSelect(layer, select, cotangent, active, residuals, cotangents),
393 .custom_call => |custom| {
394 const Rule = @typeInfo(@TypeOf(rules)).pointer.child;
395 if (comptime @hasDecl(Rule, "customCall")) {
396 var rule_ctx = CustomVjpContext(@TypeOf(layer)){
397 .layer = layer,
398 .custom = custom,
399 .cotangent = cotangent,
400 .active = active,
401 .residuals = residuals,
402 .cotangents = cotangents,
403 };
404 try rules.customCall(&rule_ctx);
405 } else {
406 return error.CustomCallRequiresVjpContract;
407 }
408 },
409 .scan, .projection => return error.ScanTransposeUnsupported,
410 }
411 }
412 }
413
414 pub fn CustomVjpContext(comptime Layer: type) type {
415 return struct {
416 layer: Layer,
417 custom: program_mod.CustomCall,
418 cotangent: trace.Value,
419 active: []const bool,
420 residuals: []const ?trace.Value,
421 cotangents: []?trace.Value,
422
423 pub fn operandCount(self: *const @This()) usize {
424 return self.custom.operands.len;
425 }
426
427 pub fn operandIsActive(self: *const @This(), index: usize) bool {
428 return self.active[self.custom.operands[index].index];
429 }
430
431 pub fn operandResidual(self: *const @This(), index: usize) trace.Value {
432 return requireResidual(self.residuals, self.custom.operands[index]);
433 }
434
435 pub fn builderHandle(self: *@This()) *trace.Builder {
436 return self.layer.builderHandle();
437 }
438
439 pub fn contribute(self: *@This(), index: usize, value: trace.Value) !void {
440 try addCotangent(self.layer, self.cotangents, self.custom.operands[index], value);
441 }
442 };
443 }
444
445 fn transposeUnary(
446 layer: anytype,
447 unary: program_mod.UnaryOp,
448 cotangent: trace.Value,
449 cotangents: []?trace.Value,
450 ) !void {
451 const contribution = switch (unary.op) {
452 .neg => try emit.unary(layer, .neg, cotangent),
453 else => return error.UnsupportedTranspose,
454 };
455 try addCotangent(layer, cotangents, unary.input, contribution);
456 }
457
458 fn transposeBinary(
459 layer: anytype,
460 binary: program_mod.BinaryOp,
461 cotangent: trace.Value,
462 active: []const bool,
463 residuals: []const ?trace.Value,
464 cotangents: []?trace.Value,
465 ) !void {
466 const lhs_active = active[binary.lhs.index];
467 const rhs_active = active[binary.rhs.index];
468
469 switch (binary.op) {
470 .add => {
471 if (lhs_active) try addCotangent(layer, cotangents, binary.lhs, cotangent);
472 if (rhs_active) try addCotangent(layer, cotangents, binary.rhs, cotangent);
473 },
474 .sub => {
475 if (lhs_active) try addCotangent(layer, cotangents, binary.lhs, cotangent);
476 if (rhs_active) try addCotangent(layer, cotangents, binary.rhs, try emit.unary(layer, .neg, cotangent));
477 },
478 .mul => {
479 if (lhs_active and rhs_active) return error.NonlinearTranspose;
480 if (lhs_active) try addCotangent(
481 layer,
482 cotangents,
483 binary.lhs,
484 try emit.binary(layer, .mul, cotangent, requireResidual(residuals, binary.rhs)),
485 );
486 if (rhs_active) try addCotangent(
487 layer,
488 cotangents,
489 binary.rhs,
490 try emit.binary(layer, .mul, requireResidual(residuals, binary.lhs), cotangent),
491 );
492 },
493 .div => {
494 if (rhs_active) return error.UnsupportedTranspose;
495 if (lhs_active) try addCotangent(
496 layer,
497 cotangents,
498 binary.lhs,
499 try emit.binary(layer, .div, cotangent, requireResidual(residuals, binary.rhs)),
500 );
501 },
502 .max, .min, .pow => return error.UnsupportedTranspose,
503 }
504 }
505
506 fn transposeSelect(
507 layer: anytype,
508 select: program_mod.Select,
509 cotangent: trace.Value,
510 active: []const bool,
511 residuals: []const ?trace.Value,
512 cotangents: []?trace.Value,
513 ) !void {
514 if (active[select.pred.index]) return error.UnsupportedTranspose;
515 const pred = requireResidual(residuals, select.pred);
516 const zeros = try emit.zeros(layer, cotangent.ty);
517 if (active[select.on_true.index]) {
518 try addCotangent(layer, cotangents, select.on_true, try emit.select(layer, pred, cotangent, zeros));
519 }
520 if (active[select.on_false.index]) {
521 try addCotangent(layer, cotangents, select.on_false, try emit.select(layer, pred, zeros, cotangent));
522 }
523 }
524
525 fn transposeBroadcast(
526 layer: anytype,
527 program: *const program_mod.Program,
528 op: *const program_mod.Operation,
529 broadcast: program_mod.Broadcast,
530 cotangent: trace.Value,
531 active: []const bool,
532 cotangents: []?trace.Value,
533 ) !void {
534 if (!active[broadcast.input.index]) return;
535 const input_ty = program.typeOf(broadcast.input);
536 if (input_ty.rank() != 0) return error.UnsupportedTranspose;
537 const axes = try allAxes(layer.builderHandle().arena.allocator(), op.result.rank());
538 const reduced = try reduceSum(layer, cotangent, axes);
539 try addCotangent(layer, cotangents, broadcast.input, reduced);
540 }
541
542 fn transposeBroadcastInDim(
543 layer: anytype,
544 program: *const program_mod.Program,
545 op: *const program_mod.Operation,
546 broadcast: program_mod.BroadcastInDim,
547 cotangent: trace.Value,
548 active: []const bool,
549 cotangents: []?trace.Value,
550 ) !void {
551 if (!active[broadcast.input.index]) return;
552 const input_ty = program.typeOf(broadcast.input);
553 const axes = try broadcastReductionAxes(
554 layer.builderHandle().arena.allocator(),
555 input_ty.dims,
556 op.result.dims,
557 broadcast.broadcast_dims,
558 );
559 var contribution = cotangent;
560 if (axes.len != 0) contribution = try reduceSum(layer, contribution, axes);
561 if (!types.sameDims(contribution.ty.dims, input_ty.dims)) {
562 contribution = try emit.reshape(layer, contribution, input_ty.dims);
563 }
564 try addCotangent(layer, cotangents, broadcast.input, contribution);
565 }
566
567 fn transposeReshape(
568 layer: anytype,
569 program: *const program_mod.Program,
570 reshape: program_mod.Reshape,
571 cotangent: trace.Value,
572 active: []const bool,
573 cotangents: []?trace.Value,
574 ) !void {
575 if (!active[reshape.input.index]) return;
576 const input_ty = program.typeOf(reshape.input);
577 try addCotangent(layer, cotangents, reshape.input, try emit.reshape(layer, cotangent, input_ty.dims));
578 }
579
580 fn transposeTranspose(
581 layer: anytype,
582 program: *const program_mod.Program,
583 transpose: program_mod.Transpose,
584 cotangent: trace.Value,
585 active: []const bool,
586 cotangents: []?trace.Value,
587 ) !void {
588 if (!active[transpose.input.index]) return;
589 const inverse = try inversePermutation(layer.builderHandle().arena.allocator(), transpose.permutation);
590 _ = program;
591 try addCotangent(layer, cotangents, transpose.input, try emit.transpose(layer, cotangent, inverse));
592 }
593
594 fn transposeReduce(
595 layer: anytype,
596 program: *const program_mod.Program,
597 reduce: program_mod.Reduce,
598 cotangent: trace.Value,
599 active: []const bool,
600 cotangents: []?trace.Value,
601 ) !void {
602 if (reduce.reducer != .sum) return error.UnsupportedTranspose;
603 if (active[reduce.init.index]) return error.UnsupportedTranspose;
604 if (!active[reduce.input.index]) return;
605 const input_ty = program.typeOf(reduce.input);
606 const dims = try keptAxes(layer.builderHandle().arena.allocator(), input_ty.rank(), reduce.dimensions);
607 const contribution = try emit.broadcastInDim(layer, cotangent, input_ty.dims, dims);
608 try addCotangent(layer, cotangents, reduce.input, contribution);
609 }
610
611 fn transposeGather(
612 layer: anytype,
613 program: *const program_mod.Program,
614 gather: program_mod.Gather,
615 cotangent: trace.Value,
616 active: []const bool,
617 residuals: []const ?trace.Value,
618 cotangents: []?trace.Value,
619 ) !void {
620 if (!active[gather.input.index]) return;
621 if (active[gather.indices.index]) return error.UnsupportedTranspose;
622 const input_ty = program.typeOf(gather.input);
623 const zeros = try emit.zeros(layer, input_ty);
624 const contribution = try emit.scatterAdd(layer, zeros, requireResidual(residuals, gather.indices), cotangent, gather.axis);
625 try addCotangent(layer, cotangents, gather.input, contribution);
626 }
627
628 fn transposeScatterAdd(
629 layer: anytype,
630 scatter_add: program_mod.ScatterAdd,
631 cotangent: trace.Value,
632 active: []const bool,
633 residuals: []const ?trace.Value,
634 cotangents: []?trace.Value,
635 ) !void {
636 if (active[scatter_add.indices.index]) return error.UnsupportedTranspose;
637 if (active[scatter_add.input.index]) {
638 try addCotangent(layer, cotangents, scatter_add.input, cotangent);
639 }
640 if (active[scatter_add.updates.index]) {
641 const contribution = try emit.gather(layer, cotangent, requireResidual(residuals, scatter_add.indices), scatter_add.axis);
642 try addCotangent(layer, cotangents, scatter_add.updates, contribution);
643 }
644 }
645
646 fn transposeDot(
647 layer: anytype,
648 program: *const program_mod.Program,
649 dot: program_mod.DotGeneral,
650 cotangent: trace.Value,
651 active: []const bool,
652 residuals: []const ?trace.Value,
653 cotangents: []?trace.Value,
654 ) !void {
655 const lhs_active = active[dot.lhs.index];
656 const rhs_active = active[dot.rhs.index];
657 if (lhs_active and rhs_active) return error.NonlinearTranspose;
658
659 const allocator = layer.builderHandle().arena.allocator();
660 const lhs_ty = program.typeOf(dot.lhs);
661 const rhs_ty = program.typeOf(dot.rhs);
662 const batch_count = dot.lhs_batch.len;
663 const lhs_free = try freeAxes(allocator, lhs_ty.rank(), dot.lhs_batch, dot.lhs_contract);
664 const rhs_free = try freeAxes(allocator, rhs_ty.rank(), dot.rhs_batch, dot.rhs_contract);
665 const cotangent_batch = try axisRange(allocator, 0, batch_count);
666 const cotangent_rows = try axisRange(allocator, batch_count, lhs_free.len);
667 const cotangent_columns = try axisRange(allocator, batch_count + lhs_free.len, rhs_free.len);
668
669 if (lhs_active) {
670 const rhs = requireResidual(residuals, dot.rhs);
671 const flat_cotangent = try collapseBlocks(layer, cotangent, cotangent_batch, cotangent_rows, cotangent_columns);
672 const flat_rhs = try collapseBlocks(layer, rhs, dot.rhs_batch, rhs_free, dot.rhs_contract);
673 const flat = try canonicalDot(layer, flat_cotangent, flat_rhs, batch_count > 0);
674 const restored = try expandBlocks(layer, flat, lhs_ty.dims, dot.lhs_batch, lhs_free, dot.lhs_contract);
675 try addCotangent(layer, cotangents, dot.lhs, restored);
676 }
677
678 if (rhs_active) {
679 const lhs = requireResidual(residuals, dot.lhs);
680 const flat_lhs = try collapseBlocks(layer, lhs, dot.lhs_batch, dot.lhs_contract, lhs_free);
681 const flat_cotangent = try collapseBlocks(layer, cotangent, cotangent_batch, cotangent_rows, cotangent_columns);
682 const flat = try canonicalDot(layer, flat_lhs, flat_cotangent, batch_count > 0);
683 const restored = try expandBlocks(layer, flat, rhs_ty.dims, dot.rhs_batch, dot.rhs_contract, rhs_free);
684 try addCotangent(layer, cotangents, dot.rhs, restored);
685 }
686 }
687
688 fn collapseBlocks(
689 layer: anytype,
690 value: trace.Value,
691 batch: []const i64,
692 rows: []const i64,
693 columns: []const i64,
694 ) !trace.Value {
695 const allocator = layer.builderHandle().arena.allocator();
696 const rank = value.ty.rank();
697 const permutation = try allocator.alloc(i64, rank);
698 var out: usize = 0;
699 for ([_][]const i64{ batch, rows, columns }) |block| {
700 for (block) |axis| {
701 permutation[out] = axis;
702 out += 1;
703 }
704 }
705 const ordered = try restoreAxisOrder(layer, value, permutation);
706
707 const collapsed_rank: usize = if (batch.len > 0) 3 else 2;
708 const collapsed = try allocator.alloc(types.Dim, collapsed_rank);
709 var dim_index: usize = 0;
710 if (batch.len > 0) {
711 collapsed[dim_index] = blockDim(ordered.ty.dims[0..batch.len], "flatbatch");
712 dim_index += 1;
713 }
714 collapsed[dim_index] = blockDim(ordered.ty.dims[batch.len .. batch.len + rows.len], "flatrows");
715 dim_index += 1;
716 collapsed[dim_index] = blockDim(ordered.ty.dims[batch.len + rows.len ..], "flatcolumns");
717
718 if (collapsed_rank == rank) {
719 var unchanged = true;
720 for (collapsed, ordered.ty.dims) |target, source| {
721 if (target.extent != source.extent) unchanged = false;
722 }
723 if (unchanged) return ordered;
724 }
725 return emit.reshape(layer, ordered, collapsed);
726 }
727
728 fn blockDim(dims: []const types.Dim, fallback: []const u8) types.Dim {
729 if (dims.len == 1) return dims[0];
730 var extent: i64 = 1;
731 for (dims) |dim| {
732 extent *= dim.extent;
733 }
734 return .{ .name = fallback, .extent = extent };
735 }
736
737 fn canonicalDot(layer: anytype, lhs: trace.Value, rhs: trace.Value, batched: bool) !trace.Value {
738 if (batched) {
739 return emit.dotGeneral(layer, lhs, rhs, &.{2}, &.{1}, &.{0}, &.{0});
740 }
741 return emit.dotGeneral(layer, lhs, rhs, &.{1}, &.{0}, &.{}, &.{});
742 }
743
744 fn expandBlocks(
745 layer: anytype,
746 flat: trace.Value,
747 target_dims: []const types.Dim,
748 batch: []const i64,
749 rows: []const i64,
750 columns: []const i64,
751 ) !trace.Value {
752 const allocator = layer.builderHandle().arena.allocator();
753 const rank = target_dims.len;
754 const expanded = try allocator.alloc(types.Dim, rank);
755 const permutation = try allocator.alloc(i64, rank);
756 var out: usize = 0;
757 for ([_][]const i64{ batch, rows, columns }) |block| {
758 for (block) |axis| {
759 expanded[out] = target_dims[@intCast(axis)];
760 permutation[@intCast(axis)] = @intCast(out);
761 out += 1;
762 }
763 }
764 var reshaped = flat;
765 if (rank != flat.ty.rank()) {
766 reshaped = try emit.reshape(layer, flat, expanded);
767 } else {
768 var unchanged = true;
769 for (expanded, flat.ty.dims) |target, source| {
770 if (target.extent != source.extent) unchanged = false;
771 }
772 if (!unchanged) reshaped = try emit.reshape(layer, flat, expanded);
773 }
774 return restoreAxisOrder(layer, reshaped, permutation);
775 }
776
777 fn restoreAxisOrder(layer: anytype, value: trace.Value, permutation: []const i64) !trace.Value {
778 for (permutation, 0..) |axis, index| {
779 if (axis != @as(i64, @intCast(index))) return emit.transpose(layer, value, permutation);
780 }
781 return value;
782 }
783
784 fn freeAxes(allocator: std.mem.Allocator, rank: usize, batch: []const i64, contract: []const i64) ![]const i64 {
785 const result = try allocator.alloc(i64, rank - batch.len - contract.len);
786 var out: usize = 0;
787 for (0..rank) |axis| {
788 if (containsAxis(batch, axis) or containsAxis(contract, axis)) continue;
789 result[out] = @intCast(axis);
790 out += 1;
791 }
792 return result;
793 }
794
795 fn axisRange(allocator: std.mem.Allocator, start: usize, count: usize) ![]const i64 {
796 const result = try allocator.alloc(i64, count);
797 for (result, 0..) |*axis, index| {
798 axis.* = @intCast(start + index);
799 }
800 return result;
801 }
802
803 fn addCotangent(
804 layer: anytype,
805 cotangents: []?trace.Value,
806 id: program_mod.Id,
807 contribution: trace.Value,
808 ) !void {
809 if (cotangents[id.index]) |existing| {
810 cotangents[id.index] = try emit.binary(layer, .add, existing, contribution);
811 } else {
812 cotangents[id.index] = contribution;
813 }
814 }
815
816 fn requireResidual(residuals: []const ?trace.Value, id: program_mod.Id) trace.Value {
817 return residuals[id.index] orelse unreachable;
818 }
819
820 fn reduceSum(layer: anytype, value: trace.Value, axes: []const i64) !trace.Value {
821 const init = try emit.fullFloat(layer, .{ .dtype = value.ty.dtype, .dims = &.{} }, 0.0);
822 return emit.reduce(layer, value, init, .sum, axes);
823 }
824
825 fn allAxes(allocator: std.mem.Allocator, rank: usize) ![]const i64 {
826 const result = try allocator.alloc(i64, rank);
827 for (result, 0..) |*axis, index| {
828 axis.* = @intCast(index);
829 }
830 return result;
831 }
832
833 fn keptAxes(allocator: std.mem.Allocator, rank: usize, removed: []const i64) ![]const i64 {
834 const result = try allocator.alloc(i64, rank - removed.len);
835 var out: usize = 0;
836 for (0..rank) |axis| {
837 if (!containsAxis(removed, axis)) {
838 result[out] = @intCast(axis);
839 out += 1;
840 }
841 }
842 return result;
843 }
844
845 fn broadcastReductionAxes(
846 allocator: std.mem.Allocator,
847 input_dims: []const types.Dim,
848 result_dims: []const types.Dim,
849 broadcast_dims: []const i64,
850 ) ![]const i64 {
851 var count: usize = 0;
852 for (0..result_dims.len) |axis| {
853 if (broadcastInputAxis(broadcast_dims, axis)) |input_axis| {
854 if (input_dims[input_axis].extent == 1 and result_dims[axis].extent != 1) count += 1;
855 } else {
856 count += 1;
857 }
858 }
859 const result = try allocator.alloc(i64, count);
860 var out: usize = 0;
861 for (0..result_dims.len) |axis| {
862 if (broadcastInputAxis(broadcast_dims, axis)) |input_axis| {
863 if (input_dims[input_axis].extent == 1 and result_dims[axis].extent != 1) {
864 result[out] = @intCast(axis);
865 out += 1;
866 }
867 } else {
868 result[out] = @intCast(axis);
869 out += 1;
870 }
871 }
872 return result;
873 }
874
875 fn broadcastInputAxis(broadcast_dims: []const i64, result_axis: usize) ?usize {
876 for (broadcast_dims, 0..) |axis, input_axis| {
877 if (axis == result_axis) return input_axis;
878 }
879 return null;
880 }
881
882 fn containsAxis(axes: []const i64, candidate: usize) bool {
883 for (axes) |axis| {
884 if (axis == candidate) return true;
885 }
886 return false;
887 }
888
889 fn inversePermutation(allocator: std.mem.Allocator, permutation: []const i64) ![]const i64 {
890 const result = try allocator.alloc(i64, permutation.len);
891 for (permutation, 0..) |axis, index| {
892 result[@intCast(axis)] = @intCast(index);
893 }
894 return result;
895 }
896
897 fn pullbackBody(_: *trace.Builder, args: []const trace.Value) !trace.Value {
898 return try (try args[0].mul(args[1])).tanh();
899 }
900
901 test "tensor pullback transposes a linearized elementwise program" {
902 var source = try trace.define(std.testing.allocator, "pullback", &.{
903 types.spec(.f32, .{ .lane = 4 }),
904 types.spec(.f32, .{ .lane = 4 }),
905 }, pullbackBody);
906 defer source.deinit();
907
908 var linearized = try autodiff.linearize(std.testing.allocator, &source, .{ .wrt = &.{ 0, 1 } });
909 defer linearized.deinit();
910
911 var transposed = try pullback(std.testing.allocator, &linearized, .{});
912 defer transposed.deinit();
913
914 try std.testing.expectEqual(@as(usize, 3), transposed.program.parameters.len);
915 try std.testing.expectEqual(@as(usize, 2), transposed.program.outputs.len);
916 try std.testing.expectEqualSlices(usize, &.{ 0, 1 }, transposed.differentiated_parameters);
917 try types.expectExtents(&.{4}, transposed.program.typeOf(transposed.program.outputs[0]));
918 try types.expectExtents(&.{4}, transposed.program.typeOf(transposed.program.outputs[1]));
919 }
920
921 fn sumBody(_: *trace.Builder, args: []const trace.Value) !trace.Value {
922 return try args[0].sum(.lane);
923 }
924
925 test "tensor pullback transposes reduce sum" {
926 var source = try trace.define(std.testing.allocator, "pullback_sum", &.{
927 types.spec(.f32, .{ .lane = 4 }),
928 }, sumBody);
929 defer source.deinit();
930
931 var linearized = try autodiff.linearize(std.testing.allocator, &source, .{ .wrt = &.{0} });
932 defer linearized.deinit();
933
934 var transposed = try pullback(std.testing.allocator, &linearized, .{});
935 defer transposed.deinit();
936
937 try std.testing.expectEqual(@as(usize, 2), transposed.program.parameters.len);
938 try types.expectExtents(&.{4}, transposed.program.typeOf(transposed.program.parameters[0]));
939 try types.expectExtents(&.{}, transposed.program.typeOf(transposed.program.parameters[1]));
940 try std.testing.expectEqual(@as(usize, 1), transposed.program.outputs.len);
941 try types.expectExtents(&.{4}, transposed.program.typeOf(transposed.program.outputs[0]));
942 }
943
944 fn productLossBody(_: *trace.Builder, args: []const trace.Value) !trace.Value {
945 const product = try args[0].mul(args[1]);
946 return try product.sum(.lane);
947 }
948
949 const ReverseGeneratedCounts = struct {
950 mul: usize = 0,
951 broadcast_in_dim: usize = 0,
952 metadata_checks: usize = 0,
953 };
954
955 const ReverseGeneratedRewrite = struct {
956 counts: *ReverseGeneratedCounts,
957
958 pub fn mul(self: *@This(), ctx: *transform.Context) !?trace.Value {
959 if (ctx.op.id.index == program_mod.synthetic_id.index) {
960 self.counts.mul += 1;
961 self.counts.metadata_checks += 1;
962 _ = ctx.isZero(0);
963 _ = ctx.constantPayload(1);
964 }
965 return null;
966 }
967
968 pub fn broadcastInDim(self: *@This(), ctx: *transform.Context) !?trace.Value {
969 if (ctx.op.id.index == program_mod.synthetic_id.index) self.counts.broadcast_in_dim += 1;
970 return null;
971 }
972 };
973
974 test "tensor pullback binds generated transpose ops through downstream semantics" {
975 var source = try trace.define(std.testing.allocator, "pullback_generated", &.{
976 types.spec(.f32, .{ .lane = 4 }),
977 types.spec(.f32, .{ .lane = 4 }),
978 }, productLossBody);
979 defer source.deinit();
980
981 var linearized = try autodiff.linearize(std.testing.allocator, &source, .{ .wrt = &.{ 0, 1 } });
982 defer linearized.deinit();
983
984 var builder = try trace.Builder.init(std.testing.allocator, source.name);
985 errdefer builder.deinit();
986
987 var counts = ReverseGeneratedCounts{};
988 const graph = interpret.Graph{ .builder = &builder };
989 const rewrite = transform.semantics(&linearized.program, graph, ReverseGeneratedRewrite{ .counts = &counts });
990 var transposed = try pullbackWith(std.testing.allocator, &linearized, .{}, rewrite);
991 defer transposed.deinit();
992
993 try std.testing.expectEqual(@as(usize, 2), counts.mul);
994 try std.testing.expectEqual(@as(usize, 1), counts.broadcast_in_dim);
995 try std.testing.expectEqual(@as(usize, 2), counts.metadata_checks);
996 try std.testing.expectEqual(@as(usize, 3), transposed.program.parameters.len);
997 try std.testing.expectEqual(@as(usize, 2), transposed.program.outputs.len);
998 }
999
1000 fn residualExpressionBody(builder: *trace.Builder, args: []const trace.Value) !trace.Value {
1001 const one = try builder.scalar(.f32, 1.0);
1002 const scale = try args[1].add(one);
1003 const product = try args[0].mul(scale);
1004 return try product.sum(.lane);
1005 }
1006
1007 test "tensor pullback replays residual expressions through interpretation" {
1008 var source = try trace.define(std.testing.allocator, "pullback_residual_expression", &.{
1009 types.spec(.f32, .{ .lane = 4 }),
1010 types.spec(.f32, .{ .lane = 4 }),
1011 }, residualExpressionBody);
1012 defer source.deinit();
1013
1014 var linearized = try autodiff.linearize(std.testing.allocator, &source, .{ .wrt = &.{0} });
1015 defer linearized.deinit();
1016
1017 var transposed = try pullback(std.testing.allocator, &linearized, .{});
1018 defer transposed.deinit();
1019
1020 var replayed_residual_expression = false;
1021 for (transposed.program.operations) |op| {
1022 switch (op.kind) {
1023 .binary => |binary| {
1024 if (binary.op == .add) replayed_residual_expression = true;
1025 },
1026 else => {},
1027 }
1028 }
1029
1030 try std.testing.expect(replayed_residual_expression);
1031 try std.testing.expectEqual(@as(usize, 3), transposed.program.parameters.len);
1032 try std.testing.expectEqual(@as(usize, 1), transposed.program.outputs.len);
1033 try types.expectExtents(&.{4}, transposed.program.typeOf(transposed.program.outputs[0]));
1034 }
1035
1036 fn matmulBody(_: *trace.Builder, args: []const trace.Value) !trace.Value {
1037 const product = try args[0].contract(args[1], .k);
1038 return try product.sum(.{ .m, .n });
1039 }
1040
1041 test "tensor pullback transposes matmul through scalar loss" {
1042 var source = try trace.define(std.testing.allocator, "pullback_matmul", &.{
1043 types.spec(.f32, .{ .m = 2, .k = 4 }),
1044 types.spec(.f32, .{ .k = 4, .n = 3 }),
1045 }, matmulBody);
1046 defer source.deinit();
1047
1048 var linearized = try autodiff.linearize(std.testing.allocator, &source, .{ .wrt = &.{ 0, 1 } });
1049 defer linearized.deinit();
1050
1051 var transposed = try pullback(std.testing.allocator, &linearized, .{});
1052 defer transposed.deinit();
1053
1054 try std.testing.expectEqual(@as(usize, 3), transposed.program.parameters.len);
1055 try std.testing.expectEqual(@as(usize, 2), transposed.program.outputs.len);
1056 try types.expectExtents(&.{ 2, 4 }, transposed.program.typeOf(transposed.program.outputs[0]));
1057 try types.expectExtents(&.{ 4, 3 }, transposed.program.typeOf(transposed.program.outputs[1]));
1058 }
1059
1060 fn batchedMatmulBody(_: *trace.Builder, args: []const trace.Value) !trace.Value {
1061 const product = try args[0].contract(args[1], .k);
1062 return try product.sum(.{ .b, .m, .n });
1063 }
1064
1065 fn rawTransposedMatmulBody(_: *trace.Builder, args: []const trace.Value) !trace.Value {
1066 const product = try args[0].builder.dotGeneralOp(args[0], args[1], &.{0}, &.{0}, &.{}, &.{});
1067 return try product.sum(.{ .m, .n });
1068 }
1069
1070 fn rawTrailingBatchMatmulBody(_: *trace.Builder, args: []const trace.Value) !trace.Value {
1071 const product = try args[0].builder.dotGeneralOp(args[0], args[1], &.{1}, &.{0}, &.{2}, &.{2});
1072 return try product.sum(.{ .b, .m, .n });
1073 }
1074
1075 test "tensor pullback transposes batched matmul through scalar loss" {
1076 var source = try trace.define(std.testing.allocator, "pullback_batched_matmul", &.{
1077 types.spec(.f32, .{ .b = 5, .m = 2, .k = 4 }),
1078 types.spec(.f32, .{ .b = 5, .k = 4, .n = 3 }),
1079 }, batchedMatmulBody);
1080 defer source.deinit();
1081
1082 var linearized = try autodiff.linearize(std.testing.allocator, &source, .{ .wrt = &.{ 0, 1 } });
1083 defer linearized.deinit();
1084
1085 var transposed = try pullback(std.testing.allocator, &linearized, .{});
1086 defer transposed.deinit();
1087
1088 try std.testing.expectEqual(@as(usize, 3), transposed.program.parameters.len);
1089 try std.testing.expectEqual(@as(usize, 2), transposed.program.outputs.len);
1090 try types.expectExtents(&.{ 5, 2, 4 }, transposed.program.typeOf(transposed.program.outputs[0]));
1091 try types.expectExtents(&.{ 5, 4, 3 }, transposed.program.typeOf(transposed.program.outputs[1]));
1092 }
1093
1094 test "tensor pullback transposes raw noncanonical rank-2 dot_general" {
1095 var source = try trace.define(std.testing.allocator, "pullback_raw_transposed_matmul", &.{
1096 types.spec(.f32, .{ .k = 4, .m = 2 }),
1097 types.spec(.f32, .{ .k = 4, .n = 3 }),
1098 }, rawTransposedMatmulBody);
1099 defer source.deinit();
1100
1101 var linearized = try autodiff.linearize(std.testing.allocator, &source, .{ .wrt = &.{ 0, 1 } });
1102 defer linearized.deinit();
1103
1104 var transposed = try pullback(std.testing.allocator, &linearized, .{});
1105 defer transposed.deinit();
1106
1107 try std.testing.expectEqual(@as(usize, 2), transposed.program.outputs.len);
1108 try types.expectExtents(&.{ 4, 2 }, transposed.program.typeOf(transposed.program.outputs[0]));
1109 try types.expectExtents(&.{ 4, 3 }, transposed.program.typeOf(transposed.program.outputs[1]));
1110 }
1111
1112 test "tensor pullback transposes raw noncanonical batched dot_general" {
1113 var source = try trace.define(std.testing.allocator, "pullback_raw_trailing_batch_matmul", &.{
1114 types.spec(.f32, .{ .m = 2, .k = 4, .b = 5 }),
1115 types.spec(.f32, .{ .k = 4, .n = 3, .b = 5 }),
1116 }, rawTrailingBatchMatmulBody);
1117 defer source.deinit();
1118
1119 var linearized = try autodiff.linearize(std.testing.allocator, &source, .{ .wrt = &.{ 0, 1 } });
1120 defer linearized.deinit();
1121
1122 var transposed = try pullback(std.testing.allocator, &linearized, .{});
1123 defer transposed.deinit();
1124
1125 try std.testing.expectEqual(@as(usize, 2), transposed.program.outputs.len);
1126 try types.expectExtents(&.{ 2, 4, 5 }, transposed.program.typeOf(transposed.program.outputs[0]));
1127 try types.expectExtents(&.{ 4, 3, 5 }, transposed.program.typeOf(transposed.program.outputs[1]));
1128 }
1129
1130 fn doublingVjpBody(builder: *trace.Builder, args: []const trace.Value) !trace.Value {
1131 _ = builder;
1132 return args[0].builder.customCall("accy.custom.double", 1, &.{args[0]}, args[0].ty);
1133 }
1134
1135 const DoublingJvpRule = struct {
1136 pub fn bind(self: *@This(), ctx: anytype) !autodiff.Dual {
1137 _ = self;
1138 switch (ctx.op.kind) {
1139 .custom_call => |custom| {
1140 if (std.mem.eql(u8, custom.target, "accy.custom.double")) {
1141 const builder = ctx.builderHandle();
1142 return .{
1143 .primal = try builder.customCall(custom.target, custom.version, &.{ctx.args[0].primal}, ctx.op.result),
1144 .tangent = try builder.customCall(custom.target, custom.version, &.{ctx.args[0].tangent}, ctx.op.result),
1145 };
1146 }
1147 },
1148 else => {},
1149 }
1150 return ctx.default();
1151 }
1152 };
1153
1154 const DoublingVjpRule = struct {
1155 applied: *usize,
1156
1157 pub fn customCall(self: *@This(), ctx: anytype) !void {
1158 self.applied.* += 1;
1159 if (!ctx.operandIsActive(0)) return;
1160 const builder = ctx.builderHandle();
1161 const contribution = try builder.customCall(ctx.custom.target, ctx.custom.version, &.{ctx.cotangent}, ctx.cotangent.ty);
1162 try ctx.contribute(0, contribution);
1163 }
1164 };
1165
1166 fn doublingLinearization(allocator: std.mem.Allocator, source: *const program_mod.Program) !autodiff.Linearization {
1167 var builder = try trace.Builder.init(allocator, source.name);
1168 errdefer builder.deinit();
1169 const graph = interpret.Graph{ .builder = &builder };
1170 const linear = autodiff.semantics(source, graph, .{ .wrt = &.{0} });
1171 return interpret.run(allocator, source, interpret.layer(autodiff.Dual, linear, DoublingJvpRule{}));
1172 }
1173
1174 test "tensor pullback rejects active custom calls without a vjp contract" {
1175 var source = try trace.define(std.testing.allocator, "vjp_opaque_custom", &.{
1176 types.spec(.f32, .{ .lane = 4 }),
1177 }, doublingVjpBody);
1178 defer source.deinit();
1179
1180 var linearized = try doublingLinearization(std.testing.allocator, &source);
1181 defer linearized.deinit();
1182
1183 try std.testing.expectError(
1184 error.CustomCallRequiresVjpContract,
1185 pullback(std.testing.allocator, &linearized, .{}),
1186 );
1187 }
1188
1189 test "tensor pullback accepts custom call vjp contracts through rules" {
1190 var source = try trace.define(std.testing.allocator, "vjp_custom_contract", &.{
1191 types.spec(.f32, .{ .lane = 4 }),
1192 }, doublingVjpBody);
1193 defer source.deinit();
1194
1195 var linearized = try doublingLinearization(std.testing.allocator, &source);
1196 defer linearized.deinit();
1197
1198 var builder = try trace.Builder.init(std.testing.allocator, linearized.program.name);
1199 errdefer builder.deinit();
1200 const graph = interpret.Graph{ .builder = &builder };
1201
1202 var applied: usize = 0;
1203 var transposed = try pullbackWithRules(
1204 std.testing.allocator,
1205 &linearized,
1206 .{},
1207 graph,
1208 DoublingVjpRule{ .applied = &applied },
1209 );
1210 defer transposed.deinit();
1211
1212 try std.testing.expectEqual(@as(usize, 1), applied);
1213
1214 var custom_calls: usize = 0;
1215 for (transposed.program.operations) |op| {
1216 switch (op.kind) {
1217 .custom_call => custom_calls += 1,
1218 else => {},
1219 }
1220 }
1221 try std.testing.expect(custom_calls >= 1);
1222 }
1223
1224 const residual_scale_target = "accy.custom.residual.scale";
1225
1226 fn residualScaleVjpBody(builder: *trace.Builder, args: []const trace.Value) !trace.Value {
1227 const scaled = try builder.customCall(residual_scale_target, 1, &.{ args[0], args[1] }, args[0].ty);
1228 return try scaled.sum(.lane);
1229 }
1230
1231 const ResidualScaleJvpRule = struct {
1232 pub fn bind(self: *@This(), ctx: anytype) !autodiff.Dual {
1233 _ = self;
1234 switch (ctx.op.kind) {
1235 .custom_call => |custom| {
1236 if (std.mem.eql(u8, custom.target, residual_scale_target)) {
1237 const builder = ctx.builderHandle();
1238 return .{
1239 .primal = try builder.customCall(custom.target, custom.version, &.{ ctx.args[0].primal, ctx.args[1].primal }, ctx.op.result),
1240 .tangent = try builder.customCall(custom.target, custom.version, &.{ ctx.args[0].tangent, ctx.args[1].primal }, ctx.op.result),
1241 };
1242 }
1243 },
1244 else => {},
1245 }
1246 return ctx.default();
1247 }
1248 };
1249
1250 const ResidualScaleVjpRule = struct {
1251 applied: *usize,
1252 residuals_used: *usize,
1253
1254 pub fn customCall(self: *@This(), ctx: anytype) !void {
1255 if (!std.mem.eql(u8, ctx.custom.target, residual_scale_target)) return error.UnexpectedCustomCall;
1256 self.applied.* += 1;
1257 if (!ctx.operandIsActive(0)) return;
1258
1259 const factor = ctx.operandResidual(1);
1260 self.residuals_used.* += 1;
1261
1262 const builder = ctx.builderHandle();
1263 const contribution = try builder.customCall(
1264 ctx.custom.target,
1265 ctx.custom.version,
1266 &.{ ctx.cotangent, factor },
1267 ctx.cotangent.ty,
1268 );
1269 try ctx.contribute(0, contribution);
1270 }
1271 };
1272
1273 fn residualScaleLinearization(allocator: std.mem.Allocator, source: *const program_mod.Program) !autodiff.Linearization {
1274 var builder = try trace.Builder.init(allocator, source.name);
1275 errdefer builder.deinit();
1276 const graph = interpret.Graph{ .builder = &builder };
1277 const linear = autodiff.semantics(source, graph, .{ .wrt = &.{0} });
1278 return interpret.run(allocator, source, interpret.layer(autodiff.Dual, linear, ResidualScaleJvpRule{}));
1279 }
1280
1281 test "tensor pullback exposes inactive custom call operands as vjp residuals" {
1282 var source = try trace.define(std.testing.allocator, "vjp_custom_operand_residual", &.{
1283 types.spec(.f32, .{ .lane = 4 }),
1284 types.spec(.f32, .{ .lane = 4 }),
1285 }, residualScaleVjpBody);
1286 defer source.deinit();
1287
1288 var linearized = try residualScaleLinearization(std.testing.allocator, &source);
1289 defer linearized.deinit();
1290
1291 var builder = try trace.Builder.init(std.testing.allocator, linearized.program.name);
1292 errdefer builder.deinit();
1293 const graph = interpret.Graph{ .builder = &builder };
1294
1295 var applied: usize = 0;
1296 var residuals_used: usize = 0;
1297 var transposed = try pullbackWithRules(
1298 std.testing.allocator,
1299 &linearized,
1300 .{},
1301 graph,
1302 ResidualScaleVjpRule{
1303 .applied = &applied,
1304 .residuals_used = &residuals_used,
1305 },
1306 );
1307 defer transposed.deinit();
1308
1309 try std.testing.expectEqual(@as(usize, 1), applied);
1310 try std.testing.expectEqual(@as(usize, 1), residuals_used);
1311 try std.testing.expectEqual(@as(usize, 3), transposed.program.parameters.len);
1312 try std.testing.expectEqual(@as(usize, 1), transposed.program.outputs.len);
1313 try types.expectExtents(&.{4}, transposed.program.typeOf(transposed.program.outputs[0]));
1314
1315 var residual_scale_calls: usize = 0;
1316 for (transposed.program.operations) |op| {
1317 switch (op.kind) {
1318 .custom_call => |custom| {
1319 if (std.mem.eql(u8, custom.target, residual_scale_target)) {
1320 residual_scale_calls += 1;
1321 try std.testing.expectEqual(@as(usize, 2), custom.operands.len);
1322 }
1323 },
1324 else => {},
1325 }
1326 }
1327 try std.testing.expectEqual(@as(usize, 1), residual_scale_calls);
1328 }
1329
1330 const stateful_inner_target = "accy.custom.stateful.inner";
1331 const stateful_outer_target = "accy.custom.stateful.outer";
1332
1333 fn statefulVjpBody(builder: *trace.Builder, args: []const trace.Value) !trace.Value {
1334 const inner = try builder.customCall(stateful_inner_target, 1, &.{args[0]}, args[0].ty);
1335 return try builder.customCall(stateful_outer_target, 1, &.{inner}, inner.ty);
1336 }
1337
1338 const StatefulJvpRule = struct {
1339 pub fn bind(self: *@This(), ctx: anytype) !autodiff.Dual {
1340 _ = self;
1341 switch (ctx.op.kind) {
1342 .custom_call => |custom| {
1343 if (std.mem.eql(u8, custom.target, stateful_inner_target) or
1344 std.mem.eql(u8, custom.target, stateful_outer_target))
1345 {
1346 const builder = ctx.builderHandle();
1347 return .{
1348 .primal = try builder.customCall(custom.target, custom.version, &.{ctx.args[0].primal}, ctx.op.result),
1349 .tangent = try builder.customCall(custom.target, custom.version, &.{ctx.args[0].tangent}, ctx.op.result),
1350 };
1351 }
1352 },
1353 else => {},
1354 }
1355 return ctx.default();
1356 }
1357 };
1358
1359 const StatefulVjpRule = struct {
1360 outer_seen: bool = false,
1361
1362 pub fn customCall(self: *@This(), ctx: anytype) !void {
1363 if (std.mem.eql(u8, ctx.custom.target, stateful_outer_target)) {
1364 self.outer_seen = true;
1365 } else if (std.mem.eql(u8, ctx.custom.target, stateful_inner_target)) {
1366 if (!self.outer_seen) return error.VjpRuleStateWasCopied;
1367 } else {
1368 return error.UnexpectedCustomCall;
1369 }
1370
1371 if (!ctx.operandIsActive(0)) return;
1372 const builder = ctx.builderHandle();
1373 const contribution = try builder.customCall(ctx.custom.target, ctx.custom.version, &.{ctx.cotangent}, ctx.cotangent.ty);
1374 try ctx.contribute(0, contribution);
1375 }
1376 };
1377
1378 fn statefulLinearization(allocator: std.mem.Allocator, source: *const program_mod.Program) !autodiff.Linearization {
1379 var builder = try trace.Builder.init(allocator, source.name);
1380 errdefer builder.deinit();
1381 const graph = interpret.Graph{ .builder = &builder };
1382 const linear = autodiff.semantics(source, graph, .{ .wrt = &.{0} });
1383 return interpret.run(allocator, source, interpret.layer(autodiff.Dual, linear, StatefulJvpRule{}));
1384 }
1385
1386 test "tensor pullback preserves local vjp rule state between custom calls" {
1387 var source = try trace.define(std.testing.allocator, "vjp_stateful_custom_contract", &.{
1388 types.spec(.f32, .{ .lane = 4 }),
1389 }, statefulVjpBody);
1390 defer source.deinit();
1391
1392 var linearized = try statefulLinearization(std.testing.allocator, &source);
1393 defer linearized.deinit();
1394
1395 var builder = try trace.Builder.init(std.testing.allocator, linearized.program.name);
1396 errdefer builder.deinit();
1397 const graph = interpret.Graph{ .builder = &builder };
1398
1399 var transposed = try pullbackWithRules(
1400 std.testing.allocator,
1401 &linearized,
1402 .{},
1403 graph,
1404 StatefulVjpRule{},
1405 );
1406 defer transposed.deinit();
1407
1408 var inner_calls: usize = 0;
1409 var outer_calls: usize = 0;
1410 for (transposed.program.operations) |op| {
1411 switch (op.kind) {
1412 .custom_call => |custom| {
1413 if (std.mem.eql(u8, custom.target, stateful_inner_target)) inner_calls += 1;
1414 if (std.mem.eql(u8, custom.target, stateful_outer_target)) outer_calls += 1;
1415 },
1416 else => {},
1417 }
1418 }
1419 try std.testing.expectEqual(@as(usize, 1), inner_calls);
1420 try std.testing.expectEqual(@as(usize, 1), outer_calls);
1421 }