tiny.accy.tensor.nn

1const std = @import("std");23const trace = @import("trace/root.zig");4const types = @import("type/root.zig");56pub fn embedding(table: trace.Value, ids: trace.Value, comptime axis: anytype) !trace.Value {7    return embeddingNamed(table, ids, comptime trace.builder.nameOf(axis));8}910pub fn embeddingNamed(table: trace.Value, ids: trace.Value, axis_name: []const u8) !trace.Value {11    return table.builder.gatherNamed(table, ids, axis_name);12}1314pub 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}1718pub 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;2122    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;2829    const ordered_targets = try builder.alignTo(targets, target_dims);30    return builder.sparseCrossEntropyOp(logits, ordered_targets, @intCast(class_axis));31}3233pub 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}3637pub 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}4748test "tensor nn embedding gathers a table by ids" {49    var builder = try trace.Builder.init(std.testing.allocator, "nn_embedding");50    defer builder.deinit();5152    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);5556    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}6162test "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();6566    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);7071    try types.expectExtents(&.{5}, losses.ty);72    try std.testing.expectEqualStrings("token", losses.ty.dims[0].name);73    try types.expectExtents(&.{}, mean.ty);74}

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