lib/accy/src/tensor/nn.zig

daab053ee43316e1809a84551d573ddd1e5bf3d2

 1 const std = @import("std");
 2 
 3 const trace = @import("trace/root.zig");
 4 const types = @import("type/root.zig");
 5 
 6 pub fn embedding(table: trace.Value, ids: trace.Value, comptime axis: anytype) !trace.Value {
 7     return embeddingNamed(table, ids, comptime trace.builder.nameOf(axis));
 8 }
 9 
10 pub fn embeddingNamed(table: trace.Value, ids: trace.Value, axis_name: []const u8) !trace.Value {
11     return table.builder.gatherNamed(table, ids, axis_name);
12 }
13 
14 pub fn sparseCrossEntropy(logits: trace.Value, targets: trace.Value, comptime class_axis: anytype) !trace.Value {
15     return sparseCrossEntropyNamed(logits, targets, comptime trace.builder.nameOf(class_axis));
16 }
17 
18 pub fn sparseCrossEntropyNamed(logits: trace.Value, targets: trace.Value, class_axis_name: []const u8) !trace.Value {
19     if (!logits.ty.dtype.isFloat()) return error.DTypeMismatch;
20     if (targets.ty.dtype != .i32) return error.DTypeMismatch;
21 
22     const builder = logits.builder;
23     const allocator = builder.arena.allocator();
24     const class_axis = types.findDim(logits.ty.dims, class_axis_name) orelse return error.AxisNotFound;
25     const class_axes = [_]i64{@intCast(class_axis)};
26     const target_dims = try types.removeAxes(allocator, logits.ty.dims, &class_axes);
27     if (targets.ty.rank() != target_dims.len) return error.ShapeMismatch;
28 
29     const ordered_targets = try builder.alignTo(targets, target_dims);
30     return builder.sparseCrossEntropyOp(logits, ordered_targets, @intCast(class_axis));
31 }
32 
33 pub fn sparseCrossEntropyMean(logits: trace.Value, targets: trace.Value, comptime class_axis: anytype) !trace.Value {
34     return sparseCrossEntropyMeanNamed(logits, targets, comptime trace.builder.nameOf(class_axis));
35 }
36 
37 pub fn sparseCrossEntropyMeanNamed(logits: trace.Value, targets: trace.Value, class_axis_name: []const u8) !trace.Value {
38     const losses = try sparseCrossEntropyNamed(logits, targets, class_axis_name);
39     if (losses.ty.rank() == 0) return losses;
40     const allocator = logits.builder.arena.allocator();
41     const axes = try allocator.alloc([]const u8, losses.ty.rank());
42     for (losses.ty.dims, axes) |dim, *slot| {
43         slot.* = dim.name;
44     }
45     return logits.builder.meanNamed(losses, axes);
46 }
47 
48 test "tensor nn embedding gathers a table by ids" {
49     var builder = try trace.Builder.init(std.testing.allocator, "nn_embedding");
50     defer builder.deinit();
51 
52     const table = try builder.input(.f32, .{ .vocab = 32, .channel = 8 });
53     const ids = try builder.input(.i32, .{ .batch = 2, .token = 5 });
54     const out = try embedding(table, ids, .vocab);
55 
56     try types.expectExtents(&.{ 2, 5, 8 }, out.ty);
57     try std.testing.expectEqualStrings("batch", out.ty.dims[0].name);
58     try std.testing.expectEqualStrings("token", out.ty.dims[1].name);
59     try std.testing.expectEqualStrings("channel", out.ty.dims[2].name);
60 }
61 
62 test "tensor nn sparse cross entropy returns target-shaped losses" {
63     var builder = try trace.Builder.init(std.testing.allocator, "nn_sparse_cross_entropy");
64     defer builder.deinit();
65 
66     const logits = try builder.input(.f32, .{ .token = 5, .vocab = 32 });
67     const targets = try builder.input(.i32, .{ .token = 5 });
68     const losses = try sparseCrossEntropy(logits, targets, .vocab);
69     const mean = try sparseCrossEntropyMean(logits, targets, .vocab);
70 
71     try types.expectExtents(&.{5}, losses.ty);
72     try std.testing.expectEqualStrings("token", losses.ty.dims[0].name);
73     try types.expectExtents(&.{}, mean.ty);
74 }