tiny.accy.tensor.dsl.output

1const std = @import("std");2const tensor = @import("../root.zig");3const gradient = tensor.gradient;4const trace = tensor.trace;5const parameter = @import("parameter.zig");67pub fn Source(comptime body: anytype) type {8    return OutputLayout(bodyReturnPayload(body));9}1011pub fn Derived(comptime SourceProgram: type, comptime derivation: anytype, comptime options: anytype) type {12    if (comptime derivation == .jvp) return Jvp(SourceProgram.Outputs);13    if (comptime derivation == .jvp_with) return Jvp(SourceProgram.Outputs);14    if (comptime derivation == .grad) return Grad(SourceProgram.Layout, options);15    if (comptime derivation == .grad_with) return Grad(SourceProgram.Layout, options.grad);16    return SourceProgram.Outputs;17}1819fn OutputLayout(comptime Result: type) type {20    if (comptime Result == trace.Value) return Positional(1);2122    return switch (@typeInfo(Result)) {23        .array => |array| Positional(array.len),24        .pointer => |pointer| Pointer(pointer),25        .@"struct" => |info| if (info.is_tuple) Positional(info.field_names.len) else Named(Result),26        else => @compileError("tensor Program bodies must return a Value, array, slice, or struct of Values"),27    };28}2930fn Pointer(comptime pointer: std.builtin.Type.Pointer) type {31    return switch (pointer.size) {32        .slice => Dynamic,33        .one => switch (@typeInfo(pointer.child)) {34            .array => |array| Positional(array.len),35            else => @compileError("tensor Program pointer outputs must point at arrays or slices of tensor Values"),36        },37        else => @compileError("tensor Program pointer outputs must point at arrays or slices of tensor Values"),38    };39}4041fn Positional(comptime output_count: usize) type {42    return struct {43        pub const count: ?usize = output_count;4445        pub fn index(comptime _: anytype) usize {46            @compileError("tensor Program outputs are positional");47        }48    };49}5051const Dynamic = struct {52    pub const count: ?usize = null;5354    pub fn index(comptime _: anytype) usize {55        @compileError("tensor Program slice outputs are dynamic and positional");56    }57};5859fn Named(comptime Result: type) type {60    const info = @typeInfo(Result).@"struct";61    return struct {62        pub const count: ?usize = info.field_names.len;6364        pub fn index(comptime spec: anytype) usize {65            return fieldIndex(info.field_names, name(spec));66        }67    };68}6970fn Jvp(comptime SourceOutputLayout: type) type {71    return struct {72        pub const count: ?usize = if (SourceOutputLayout.count) |source_count| source_count * 2 else null;7374        pub fn index(comptime spec: anytype) usize {75            const source_count = SourceOutputLayout.count orelse @compileError("tensor jvp output lookup requires statically known source outputs");76            return switch (@typeInfo(@TypeOf(spec))) {77                .@"struct" => |info| blk: {78                    if (info.is_tuple or info.field_names.len != 1) @compileError("tensor jvp output lookup requires .primals or .tangents");79                    const field_name = info.field_names[0];80                    if (comptime std.mem.eql(u8, field_name, "primals")) {81                        break :blk SourceOutputLayout.index(@field(spec, field_name));82                    }83                    if (comptime std.mem.eql(u8, field_name, "tangents")) {84                        break :blk source_count + SourceOutputLayout.index(@field(spec, field_name));85                    }86                    @compileError("tensor jvp output lookup requires .primals or .tangents");87                },88                .enum_literal => blk: {89                    const output_name = @tagName(spec);90                    if (comptime std.mem.eql(u8, output_name, "primal")) {91                        if (source_count != 1) @compileError("tensor jvp .primal output lookup requires one source output");92                        break :blk 0;93                    }94                    if (comptime std.mem.eql(u8, output_name, "tangent")) {95                        if (source_count != 1) @compileError("tensor jvp .tangent output lookup requires one source output");96                        break :blk source_count;97                    }98                    @compileError("tensor jvp output lookup requires .primal, .tangent, .primals, or .tangents");99                },100                else => @compileError("tensor jvp output lookup requires .primal, .tangent, .primals, or .tangents"),101            };102        }103    };104}105106fn Grad(comptime SourceLayout: type, comptime options: gradient.Options) type {107    return struct {108        pub const count: ?usize = options.wrt.len;109110        pub fn index(comptime spec: anytype) usize {111            const parameter_index = comptime SourceLayout.index(spec);112            return parameter.wrtOffset(options.wrt, parameter_index);113        }114    };115}116117fn bodyReturnPayload(comptime body: anytype) type {118    const return_type = @typeInfo(@TypeOf(body)).@"fn".return_type orelse @compileError("tensor Program body must return values");119    return switch (@typeInfo(return_type)) {120        .error_union => |info| info.payload,121        else => return_type,122    };123}124125fn fieldIndex(comptime field_names: []const [:0]const u8, comptime output_name: []const u8) usize {126    inline for (field_names, 0..) |field_name, index_value| {127        if (comptime std.mem.eql(u8, field_name, output_name)) return index_value;128    }129    @compileError("unknown tensor output: " ++ output_name);130}131132fn name(comptime name_value: anytype) []const u8 {133    return switch (@typeInfo(@TypeOf(name_value))) {134        .enum_literal => @tagName(name_value),135        .pointer => |pointer| blk: {136            if (pointer.child == u8) break :blk name_value[0..];137            if (@typeInfo(pointer.child) == .array) {138                const array = @typeInfo(pointer.child).array;139                if (array.child == u8) break :blk name_value[0..array.len];140            }141            @compileError("tensor output name pointers must point at string literals");142        },143        .array => |array| blk: {144            if (array.child == u8) break :blk name_value[0..];145            @compileError("tensor output name arrays must contain bytes");146        },147        else => @compileError("tensor output names must be enum literals or string literals"),148    };149}

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