tiny.accy.choir.semantics

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

1const std = @import("std");2const choir_abi = @import("choir_abi");3const testing = std.testing;4const axis_roles = @import("../axis/root.zig").roles;5const activation = @import("activation.zig");6const einsum = @import("einsum/root.zig");78pub const Allocator = std.mem.Allocator;9pub const DType = choir_abi.DType;1011pub const Type = struct {12    dtype: DType,13    dims: []const i64,1415    pub fn eql(a: Type, b: Type) bool {16        return a.dtype == b.dtype and std.mem.eql(i64, a.dims, b.dims);17    }18};1920pub const OpKind = enum(u8) {21    constant,22    parameter,23    iota,24    add,25    sub,26    mul,27    div,28    max,29    min,30    pow,31    compare,32    neg,33    exp,34    log,35    tanh,36    sqrt,37    activation,38    abs,39    sin,40    cos,41    tan,42    floor,43    round,44    trunc,45    atan2,46    convert,47    reduce,48    dot_general,49    einsum,50    broadcast,51    broadcast_in_dim,52    reshape,53    transpose,54    slice,55    gather,56    scatter,57    scatter_add,58    sparse_cross_entropy,59    pad,60    concatenate,61    select,62    cumsum,63    scratch,64    kernel_call,65    iterate,66    iterate_yield,67    @"return",68};6970pub const CompareDirection = enum { eq, ne, lt, le, gt, ge };7172pub const ActivationKind = activation.Kind;7374pub const ReducerKind = enum { sum, max, min };7576pub const KernelOperandEffect = enum(u8) {77    none,78    read,79    write,80    read_write,81    unknown,8283    pub fn fromByte(value: u8) ?KernelOperandEffect {84        return switch (value) {85            @backingInt(KernelOperandEffect.none) => .none,86            @backingInt(KernelOperandEffect.read) => .read,87            @backingInt(KernelOperandEffect.write) => .write,88            @backingInt(KernelOperandEffect.read_write) => .read_write,89            @backingInt(KernelOperandEffect.unknown) => .unknown,90            else => null,91        };92    }9394    pub fn writes(self: KernelOperandEffect) bool {95        return switch (self) {96            .write, .read_write, .unknown => true,97            .none, .read => false,98        };99    }100};101102pub const KernelCallContract = struct {103    target: []const u8,104    version: u32,105    has_side_effects: bool,106    operand_effects: []const KernelOperandEffect,107    result_aliases: []const ?usize,108    results: []const Type,109};110111pub const Attribute = union(enum) {112    i64: i64,113    i64_list: []const i64,114    dtype: DType,115    compare_direction: CompareDirection,116    activation_kind: ActivationKind,117    reducer_kind: ReducerKind,118    kernel_call: KernelCallContract,119    bytes: []const u8,120    einsum: []const u8,121};122123pub const InferError = error{124    ArityMismatch,125    DTypeMismatch,126    ShapeMismatch,127    RankMismatch,128    AttributeMissing,129    AttributeKindMismatch,130    InvalidDimension,131    ElementCountMismatch,132    InvalidPermutation,133    DimMismatch,134    ContractionMismatch,135    InvalidEinsum,136    InvalidKernelContract,137} || Allocator.Error;138139pub const InferFn = *const fn (140    arena: Allocator,141    inputs: []const Type,142    attrs: []const Attribute,143) InferError![]const Type;144145pub const SupportedDTypes = enum {146    any,147    arithmetic,148    float_only,149150    pub fn allows(self: SupportedDTypes, dt: DType) bool {151        return switch (self) {152            .any => true,153            .arithmetic => dt.isNumeric(),154            .float_only => dt.isFloat(),155        };156    }157};158159pub const OpInfo = struct {160    kind: OpKind,161    name: []const u8,162    supported_dtypes: SupportedDTypes,163    infer: InferFn,164};165166pub fn info(kind: OpKind) OpInfo {167    return switch (kind) {168        .constant => .{ .kind = .constant, .name = "constant", .supported_dtypes = .any, .infer = inferConstant },169        .parameter => .{ .kind = .parameter, .name = "parameter", .supported_dtypes = .any, .infer = inferParameter },170        .iota => .{ .kind = .iota, .name = "iota", .supported_dtypes = .any, .infer = inferIota },171        .add => .{ .kind = .add, .name = "add", .supported_dtypes = .arithmetic, .infer = inferElementwiseBinary },172        .sub => .{ .kind = .sub, .name = "sub", .supported_dtypes = .arithmetic, .infer = inferElementwiseBinary },173        .mul => .{ .kind = .mul, .name = "mul", .supported_dtypes = .arithmetic, .infer = inferElementwiseBinary },174        .div => .{ .kind = .div, .name = "div", .supported_dtypes = .arithmetic, .infer = inferElementwiseBinary },175        .max => .{ .kind = .max, .name = "max", .supported_dtypes = .arithmetic, .infer = inferElementwiseBinary },176        .min => .{ .kind = .min, .name = "min", .supported_dtypes = .arithmetic, .infer = inferElementwiseBinary },177        .pow => .{ .kind = .pow, .name = "pow", .supported_dtypes = .float_only, .infer = inferElementwiseBinary },178        .compare => .{ .kind = .compare, .name = "compare", .supported_dtypes = .any, .infer = inferCompare },179        .neg => .{ .kind = .neg, .name = "neg", .supported_dtypes = .arithmetic, .infer = inferElementwiseUnary },180        .exp => .{ .kind = .exp, .name = "exp", .supported_dtypes = .float_only, .infer = inferElementwiseUnary },181        .log => .{ .kind = .log, .name = "log", .supported_dtypes = .float_only, .infer = inferElementwiseUnary },182        .tanh => .{ .kind = .tanh, .name = "tanh", .supported_dtypes = .float_only, .infer = inferElementwiseUnary },183        .sqrt => .{ .kind = .sqrt, .name = "sqrt", .supported_dtypes = .float_only, .infer = inferElementwiseUnary },184        .activation => .{ .kind = .activation, .name = "activation", .supported_dtypes = .float_only, .infer = inferActivation },185        .abs => .{ .kind = .abs, .name = "abs", .supported_dtypes = .arithmetic, .infer = inferElementwiseUnary },186        .sin => .{ .kind = .sin, .name = "sin", .supported_dtypes = .float_only, .infer = inferElementwiseUnary },187        .cos => .{ .kind = .cos, .name = "cos", .supported_dtypes = .float_only, .infer = inferElementwiseUnary },188        .tan => .{ .kind = .tan, .name = "tan", .supported_dtypes = .float_only, .infer = inferElementwiseUnary },189        .floor => .{ .kind = .floor, .name = "floor", .supported_dtypes = .float_only, .infer = inferElementwiseUnary },190        .round => .{ .kind = .round, .name = "round", .supported_dtypes = .float_only, .infer = inferElementwiseUnary },191        .trunc => .{ .kind = .trunc, .name = "trunc", .supported_dtypes = .float_only, .infer = inferElementwiseUnary },192        .atan2 => .{ .kind = .atan2, .name = "atan2", .supported_dtypes = .float_only, .infer = inferElementwiseBinary },193        .convert => .{ .kind = .convert, .name = "convert", .supported_dtypes = .any, .infer = inferConvert },194        .reduce => .{ .kind = .reduce, .name = "reduce", .supported_dtypes = .any, .infer = inferReduce },195        .dot_general => .{ .kind = .dot_general, .name = "dot_general", .supported_dtypes = .arithmetic, .infer = inferDotGeneral },196        .einsum => .{ .kind = .einsum, .name = "einsum", .supported_dtypes = .arithmetic, .infer = inferEinsum },197        .broadcast => .{ .kind = .broadcast, .name = "broadcast", .supported_dtypes = .any, .infer = inferBroadcast },198        .broadcast_in_dim => .{ .kind = .broadcast_in_dim, .name = "broadcast_in_dim", .supported_dtypes = .any, .infer = inferBroadcastInDim },199        .reshape => .{ .kind = .reshape, .name = "reshape", .supported_dtypes = .any, .infer = inferReshape },200        .transpose => .{ .kind = .transpose, .name = "transpose", .supported_dtypes = .any, .infer = inferTranspose },201        .slice => .{ .kind = .slice, .name = "slice", .supported_dtypes = .any, .infer = inferSlice },202        .gather => .{ .kind = .gather, .name = "gather", .supported_dtypes = .any, .infer = inferGather },203        .scatter => .{ .kind = .scatter, .name = "scatter", .supported_dtypes = .any, .infer = inferScatter },204        .scatter_add => .{ .kind = .scatter_add, .name = "scatter_add", .supported_dtypes = .any, .infer = inferScatter },205        .sparse_cross_entropy => .{ .kind = .sparse_cross_entropy, .name = "sparse_cross_entropy", .supported_dtypes = .any, .infer = inferSparseCrossEntropy },206        .pad => .{ .kind = .pad, .name = "pad", .supported_dtypes = .any, .infer = inferPad },207        .concatenate => .{ .kind = .concatenate, .name = "concatenate", .supported_dtypes = .any, .infer = inferConcatenate },208        .select => .{ .kind = .select, .name = "select", .supported_dtypes = .any, .infer = inferSelect },209        .cumsum => .{ .kind = .cumsum, .name = "cumsum", .supported_dtypes = .arithmetic, .infer = inferCumsum },210        .scratch => .{ .kind = .scratch, .name = "scratch", .supported_dtypes = .any, .infer = inferScratch },211        .kernel_call => .{ .kind = .kernel_call, .name = "kernel_call", .supported_dtypes = .any, .infer = inferKernelCall },212        .iterate => .{ .kind = .iterate, .name = "iterate", .supported_dtypes = .any, .infer = inferIterate },213        .iterate_yield => .{ .kind = .iterate_yield, .name = "iterate_yield", .supported_dtypes = .any, .infer = inferIterateYield },214        .@"return" => .{ .kind = .@"return", .name = "return", .supported_dtypes = .any, .infer = inferReturn },215    };216}217218pub fn inferShape(219    kind: OpKind,220    arena: Allocator,221    inputs: []const Type,222    attrs: []const Attribute,223) InferError![]const Type {224    return try info(kind).infer(arena, inputs, attrs);225}226227fn inferAttrI64(attrs: []const Attribute, idx: usize) InferError!i64 {228    if (idx >= attrs.len) return error.AttributeMissing;229    return switch (attrs[idx]) {230        .i64 => |v| v,231        else => error.AttributeKindMismatch,232    };233}234235fn inferAttrI64List(attrs: []const Attribute, idx: usize) InferError![]const i64 {236    if (idx >= attrs.len) return error.AttributeMissing;237    return switch (attrs[idx]) {238        .i64_list => |v| v,239        else => error.AttributeKindMismatch,240    };241}242243fn inferAttrDType(attrs: []const Attribute, idx: usize) InferError!DType {244    if (idx >= attrs.len) return error.AttributeMissing;245    return switch (attrs[idx]) {246        .dtype => |v| v,247        else => error.AttributeKindMismatch,248    };249}250251fn inferAttrCompareDirection(attrs: []const Attribute, idx: usize) InferError!CompareDirection {252    if (idx >= attrs.len) return error.AttributeMissing;253    return switch (attrs[idx]) {254        .compare_direction => |v| v,255        else => error.AttributeKindMismatch,256    };257}258259fn inferAttrActivationKind(attrs: []const Attribute, idx: usize) InferError!ActivationKind {260    if (idx >= attrs.len) return error.AttributeMissing;261    return switch (attrs[idx]) {262        .activation_kind => |v| v,263        else => error.AttributeKindMismatch,264    };265}266267fn inferAttrReducerKind(attrs: []const Attribute, idx: usize) InferError!ReducerKind {268    if (idx >= attrs.len) return error.AttributeMissing;269    return switch (attrs[idx]) {270        .reducer_kind => |v| v,271        else => error.AttributeKindMismatch,272    };273}274275fn inferAttrBytes(attrs: []const Attribute, idx: usize) InferError![]const u8 {276    if (idx >= attrs.len) return error.AttributeMissing;277    return switch (attrs[idx]) {278        .bytes => |v| v,279        else => error.AttributeKindMismatch,280    };281}282283fn inferAttrEinsum(attrs: []const Attribute, idx: usize) InferError![]const u8 {284    if (idx >= attrs.len) return error.AttributeMissing;285    return switch (attrs[idx]) {286        .einsum => |v| v,287        else => error.AttributeKindMismatch,288    };289}290291fn inferAttrKernelCall(attrs: []const Attribute, idx: usize) InferError!KernelCallContract {292    if (idx >= attrs.len) return error.AttributeMissing;293    return switch (attrs[idx]) {294        .kernel_call => |v| v,295        else => error.AttributeKindMismatch,296    };297}298299fn singletonType(arena: Allocator, t: Type) InferError![]const Type {300    const out = try arena.alloc(Type, 1);301    out[0] = t;302    return out;303}304305fn duplicateDims(arena: Allocator, dims: []const i64) InferError![]const i64 {306    return arena.dupe(i64, dims);307}308309fn duplicateTypes(arena: Allocator, types: []const Type) InferError![]const Type {310    const out = try arena.alloc(Type, types.len);311    for (types, 0..) |typ, i| {312        for (typ.dims) |d| if (d < 0) return error.InvalidDimension;313        out[i] = .{314            .dtype = typ.dtype,315            .dims = try duplicateDims(arena, typ.dims),316        };317    }318    return out;319}320321pub fn inferConstant(arena: Allocator, inputs: []const Type, attrs: []const Attribute) InferError![]const Type {322    if (inputs.len != 0) return error.ArityMismatch;323    _ = try inferAttrBytes(attrs, 0);324    const dt = try inferAttrDType(attrs, 1);325    const dims = try inferAttrI64List(attrs, 2);326    return singletonType(arena, .{327        .dtype = dt,328        .dims = try duplicateDims(arena, dims),329    });330}331332pub fn inferParameter(arena: Allocator, inputs: []const Type, attrs: []const Attribute) InferError![]const Type {333    if (inputs.len != 0) return error.ArityMismatch;334    const idx = try inferAttrI64(attrs, 0);335    if (idx < 0) return error.InvalidDimension;336    const dt = try inferAttrDType(attrs, 1);337    const dims = try inferAttrI64List(attrs, 2);338    return singletonType(arena, .{339        .dtype = dt,340        .dims = try duplicateDims(arena, dims),341    });342}343344pub fn inferIota(arena: Allocator, inputs: []const Type, attrs: []const Attribute) InferError![]const Type {345    if (inputs.len != 0) return error.ArityMismatch;346    const dim = try inferAttrI64(attrs, 0);347    const dt = try inferAttrDType(attrs, 1);348    const dims = try inferAttrI64List(attrs, 2);349    if (dim < 0 or dim >= @as(i64, @intCast(dims.len))) return error.InvalidDimension;350    return singletonType(arena, .{351        .dtype = dt,352        .dims = try duplicateDims(arena, dims),353    });354}355356pub fn inferElementwiseBinary(arena: Allocator, inputs: []const Type, attrs: []const Attribute) InferError![]const Type {357    _ = attrs;358    if (inputs.len != 2) return error.ArityMismatch;359    if (inputs[0].dtype != inputs[1].dtype) return error.DTypeMismatch;360    if (!std.mem.eql(i64, inputs[0].dims, inputs[1].dims)) return error.ShapeMismatch;361    return singletonType(arena, .{362        .dtype = inputs[0].dtype,363        .dims = try duplicateDims(arena, inputs[0].dims),364    });365}366367pub fn inferCompare(arena: Allocator, inputs: []const Type, attrs: []const Attribute) InferError![]const Type {368    _ = try inferAttrCompareDirection(attrs, 0);369    if (inputs.len != 2) return error.ArityMismatch;370    if (inputs[0].dtype != inputs[1].dtype) return error.DTypeMismatch;371    if (!std.mem.eql(i64, inputs[0].dims, inputs[1].dims)) return error.ShapeMismatch;372    return singletonType(arena, .{373        .dtype = .i1,374        .dims = try duplicateDims(arena, inputs[0].dims),375    });376}377378pub fn inferElementwiseUnary(arena: Allocator, inputs: []const Type, attrs: []const Attribute) InferError![]const Type {379    _ = attrs;380    if (inputs.len != 1) return error.ArityMismatch;381    return singletonType(arena, .{382        .dtype = inputs[0].dtype,383        .dims = try duplicateDims(arena, inputs[0].dims),384    });385}386387pub fn inferActivation(arena: Allocator, inputs: []const Type, attrs: []const Attribute) InferError![]const Type {388    _ = try inferAttrActivationKind(attrs, 0);389    return inferElementwiseUnary(arena, inputs, attrs);390}391392pub fn inferConvert(arena: Allocator, inputs: []const Type, attrs: []const Attribute) InferError![]const Type {393    if (inputs.len != 1) return error.ArityMismatch;394    const target = try inferAttrDType(attrs, 0);395    return singletonType(arena, .{396        .dtype = target,397        .dims = try duplicateDims(arena, inputs[0].dims),398    });399}400401pub fn inferReduce(arena: Allocator, inputs: []const Type, attrs: []const Attribute) InferError![]const Type {402    if (inputs.len != 2) return error.ArityMismatch;403    _ = try inferAttrReducerKind(attrs, 0);404    const dims = try inferAttrI64List(attrs, 1);405    const operand = inputs[0];406    const init = inputs[1];407    if (operand.dtype != init.dtype) return error.DTypeMismatch;408    if (init.dims.len != 0) return error.RankMismatch;409    const rank: i64 = @intCast(operand.dims.len);410    for (dims, 0..) |d, i| {411        if (d < 0 or d >= rank) return error.InvalidDimension;412        for (dims[0..i]) |e| if (e == d) return error.InvalidDimension;413    }414    const out_rank = operand.dims.len - dims.len;415    const out_dims = try arena.alloc(i64, out_rank);416    var j: usize = 0;417    outer: for (operand.dims, 0..) |size, idx_u| {418        const idx: i64 = @intCast(idx_u);419        for (dims) |d| if (d == idx) continue :outer;420        out_dims[j] = size;421        j += 1;422    }423    std.debug.assert(j == out_rank);424    return singletonType(arena, .{ .dtype = operand.dtype, .dims = out_dims });425}426427pub fn inferDotGeneral(arena: Allocator, inputs: []const Type, attrs: []const Attribute) InferError![]const Type {428    if (inputs.len != 2) return error.ArityMismatch;429    const lhs = inputs[0];430    const rhs = inputs[1];431    if (lhs.dtype != rhs.dtype) return error.DTypeMismatch;432    const lhs_batch = try inferAttrI64List(attrs, 0);433    const rhs_batch = try inferAttrI64List(attrs, 1);434    const lhs_contract = try inferAttrI64List(attrs, 2);435    const rhs_contract = try inferAttrI64List(attrs, 3);436    const result_dtype = try inferDotGeneralResultDType(lhs.dtype, attrs);437    if (lhs_batch.len != rhs_batch.len) return error.ContractionMismatch;438    if (lhs_contract.len != rhs_contract.len) return error.ContractionMismatch;439440    if (axis_roles.check(lhs.dims.len, &.{ lhs_batch, lhs_contract }) != null) {441        return error.InvalidDimension;442    }443    if (axis_roles.check(rhs.dims.len, &.{ rhs_batch, rhs_contract }) != null) {444        return error.InvalidDimension;445    }446447    for (lhs_batch, rhs_batch) |ld, rd| {448        if (lhs.dims[@intCast(ld)] != rhs.dims[@intCast(rd)]) return error.ContractionMismatch;449    }450    for (lhs_contract, rhs_contract) |ld, rd| {451        if (lhs.dims[@intCast(ld)] != rhs.dims[@intCast(rd)]) return error.ContractionMismatch;452    }453454    const out_rank = lhs_batch.len +455        (lhs.dims.len - lhs_batch.len - lhs_contract.len) +456        (rhs.dims.len - rhs_batch.len - rhs_contract.len);457    const out_dims = try arena.alloc(i64, out_rank);458    var w: usize = 0;459    for (lhs_batch) |d| {460        out_dims[w] = lhs.dims[@intCast(d)];461        w += 1;462    }463    for (lhs.dims, 0..) |size, idx_u| {464        const idx: i64 = @intCast(idx_u);465        if (containsDim(lhs_batch, idx)) continue;466        if (containsDim(lhs_contract, idx)) continue;467        out_dims[w] = size;468        w += 1;469    }470    for (rhs.dims, 0..) |size, idx_u| {471        const idx: i64 = @intCast(idx_u);472        if (containsDim(rhs_batch, idx)) continue;473        if (containsDim(rhs_contract, idx)) continue;474        out_dims[w] = size;475        w += 1;476    }477    std.debug.assert(w == out_rank);478    return singletonType(arena, .{ .dtype = result_dtype, .dims = out_dims });479}480481pub fn inferEinsum(arena: Allocator, inputs: []const Type, attrs: []const Attribute) InferError![]const Type {482    if (inputs.len == 0) return error.ArityMismatch;483    const equation_text = try inferAttrEinsum(attrs, 0);484    const dtype = inputs[0].dtype;485    const shapes = try arena.alloc([]const u64, inputs.len);486    for (inputs, 0..) |input, input_index| {487        if (input.dtype != dtype) return error.DTypeMismatch;488        const dims = try arena.alloc(u64, input.dims.len);489        for (input.dims, 0..) |dim, dim_index| {490            if (dim < 0) return error.InvalidDimension;491            dims[dim_index] = @intCast(dim);492        }493        shapes[input_index] = dims;494    }495496    var parsed = einsum.parse(arena, equation_text, shapes) catch return error.InvalidEinsum;497    defer parsed.deinit();498499    const out_dims = try arena.alloc(i64, parsed.output.len);500    for (parsed.output, 0..) |index, i| {501        const dim = parsed.dimension(index);502        if (dim > @as(u64, @intCast(std.math.maxInt(i64)))) return error.InvalidDimension;503        out_dims[i] = @intCast(dim);504    }505    return singletonType(arena, .{ .dtype = dtype, .dims = out_dims });506}507508fn inferDotGeneralResultDType(input_dtype: DType, attrs: []const Attribute) InferError!DType {509    if (attrs.len == 4) return input_dtype;510    if (attrs.len != 5) return error.AttributeKindMismatch;511    const result_dtype = try inferAttrDType(attrs, 4);512    if (!dotGeneralResultDTypeAllowed(input_dtype, result_dtype)) return error.DTypeMismatch;513    return result_dtype;514}515516fn dotGeneralResultDTypeAllowed(input_dtype: DType, result_dtype: DType) bool {517    return input_dtype == result_dtype or (input_dtype == .f16 and result_dtype == .f32);518}519520fn containsDim(list: []const i64, d: i64) bool {521    for (list) |e| if (e == d) return true;522    return false;523}524525pub fn inferBroadcast(arena: Allocator, inputs: []const Type, attrs: []const Attribute) InferError![]const Type {526    if (inputs.len != 1) return error.ArityMismatch;527    const sizes = try inferAttrI64List(attrs, 0);528    for (sizes) |s| if (s < 0) return error.InvalidDimension;529    const out_dims = try arena.alloc(i64, sizes.len + inputs[0].dims.len);530    @memcpy(out_dims[0..sizes.len], sizes);531    @memcpy(out_dims[sizes.len..], inputs[0].dims);532    return singletonType(arena, .{ .dtype = inputs[0].dtype, .dims = out_dims });533}534535pub fn inferBroadcastInDim(arena: Allocator, inputs: []const Type, attrs: []const Attribute) InferError![]const Type {536    if (inputs.len != 1) return error.ArityMismatch;537    const broadcast_dims = try inferAttrI64List(attrs, 0);538    const result_shape = try inferAttrI64List(attrs, 1);539    const operand = inputs[0];540541    if (broadcast_dims.len != operand.dims.len) return error.ArityMismatch;542    const result_rank: i64 = @intCast(result_shape.len);543    for (broadcast_dims, 0..) |d, i| {544        if (d < 0 or d >= result_rank) return error.InvalidDimension;545        for (broadcast_dims[0..i]) |e| if (e == d) return error.InvalidDimension;546    }547    for (broadcast_dims, operand.dims) |target, operand_dim| {548        const out_dim = result_shape[@intCast(target)];549        if (operand_dim != 1 and operand_dim != out_dim) return error.ShapeMismatch;550    }551    for (result_shape) |s| if (s < 0) return error.InvalidDimension;552    return singletonType(arena, .{553        .dtype = operand.dtype,554        .dims = try duplicateDims(arena, result_shape),555    });556}557558pub fn inferReshape(arena: Allocator, inputs: []const Type, attrs: []const Attribute) InferError![]const Type {559    if (inputs.len != 1) return error.ArityMismatch;560    const new_shape = try inferAttrI64List(attrs, 0);561    for (new_shape) |s| if (s < 0) return error.InvalidDimension;562    const old_count = productI64(inputs[0].dims);563    const new_count = productI64(new_shape);564    if (old_count != new_count) return error.ElementCountMismatch;565    return singletonType(arena, .{566        .dtype = inputs[0].dtype,567        .dims = try duplicateDims(arena, new_shape),568    });569}570571fn productI64(dims: []const i64) i64 {572    var p: i64 = 1;573    for (dims) |d| p *= d;574    return p;575}576577pub fn inferTranspose(arena: Allocator, inputs: []const Type, attrs: []const Attribute) InferError![]const Type {578    if (inputs.len != 1) return error.ArityMismatch;579    const perm = try inferAttrI64List(attrs, 0);580    const operand = inputs[0];581    if (perm.len != operand.dims.len) return error.InvalidPermutation;582    const rank: i64 = @intCast(operand.dims.len);583    for (perm, 0..) |d, i| {584        if (d < 0 or d >= rank) return error.InvalidPermutation;585        for (perm[0..i]) |e| if (e == d) return error.InvalidPermutation;586    }587    const out_dims = try arena.alloc(i64, operand.dims.len);588    for (perm, 0..) |d, i| out_dims[i] = operand.dims[@intCast(d)];589    return singletonType(arena, .{ .dtype = operand.dtype, .dims = out_dims });590}591592pub fn inferSlice(arena: Allocator, inputs: []const Type, attrs: []const Attribute) InferError![]const Type {593    if (inputs.len != 1) return error.ArityMismatch;594    const start = try inferAttrI64List(attrs, 0);595    const limit = try inferAttrI64List(attrs, 1);596    const stride = try inferAttrI64List(attrs, 2);597    const operand = inputs[0];598    if (start.len != operand.dims.len or599        limit.len != operand.dims.len or600        stride.len != operand.dims.len) return error.InvalidDimension;601    const out_dims = try arena.alloc(i64, operand.dims.len);602    for (start, limit, stride, operand.dims, 0..) |s, l, st, full, i| {603        if (st <= 0) return error.InvalidDimension;604        if (s < 0 or l < s or l > full) return error.InvalidDimension;605        const span = l - s;606        out_dims[i] = @divTrunc(span + st - 1, st);607    }608    return singletonType(arena, .{ .dtype = operand.dtype, .dims = out_dims });609}610611pub fn inferGather(arena: Allocator, inputs: []const Type, attrs: []const Attribute) InferError![]const Type {612    if (inputs.len != 2) return error.ArityMismatch;613    const axis_i64 = try inferAttrI64(attrs, 0);614    const operand = inputs[0];615    const indices = inputs[1];616    if (!isIndexDType(indices.dtype)) return error.DTypeMismatch;617    const rank: i64 = @intCast(operand.dims.len);618    if (axis_i64 < 0 or axis_i64 >= rank) return error.InvalidDimension;619    const axis: usize = @intCast(axis_i64);620621    const out_rank = operand.dims.len - 1 + indices.dims.len;622    const out_dims = try arena.alloc(i64, out_rank);623    var out_i: usize = 0;624    for (operand.dims[0..axis]) |dim| {625        out_dims[out_i] = dim;626        out_i += 1;627    }628    for (indices.dims) |dim| {629        if (dim < 0) return error.InvalidDimension;630        out_dims[out_i] = dim;631        out_i += 1;632    }633    for (operand.dims[axis + 1 ..]) |dim| {634        out_dims[out_i] = dim;635        out_i += 1;636    }637    return singletonType(arena, .{ .dtype = operand.dtype, .dims = out_dims });638}639640pub fn inferScatter(arena: Allocator, inputs: []const Type, attrs: []const Attribute) InferError![]const Type {641    if (inputs.len != 3) return error.ArityMismatch;642    const axis_i64 = try inferAttrI64(attrs, 0);643    const operand = inputs[0];644    const indices = inputs[1];645    const updates = inputs[2];646    if (!isIndexDType(indices.dtype)) return error.DTypeMismatch;647    if (updates.dtype != operand.dtype) return error.DTypeMismatch;648    const rank: i64 = @intCast(operand.dims.len);649    if (axis_i64 < 0 or axis_i64 >= rank) return error.InvalidDimension;650    const axis: usize = @intCast(axis_i64);651652    const expected_rank = operand.dims.len - 1 + indices.dims.len;653    if (updates.dims.len != expected_rank) return error.RankMismatch;654655    var update_axis: usize = 0;656    for (operand.dims[0..axis]) |dim| {657        if (updates.dims[update_axis] != dim) return error.ShapeMismatch;658        update_axis += 1;659    }660    for (indices.dims) |dim| {661        if (dim < 0) return error.InvalidDimension;662        if (updates.dims[update_axis] != dim) return error.ShapeMismatch;663        update_axis += 1;664    }665    for (operand.dims[axis + 1 ..]) |dim| {666        if (updates.dims[update_axis] != dim) return error.ShapeMismatch;667        update_axis += 1;668    }669670    const out_dims = try arena.dupe(i64, operand.dims);671    return singletonType(arena, .{ .dtype = operand.dtype, .dims = out_dims });672}673674pub fn inferSparseCrossEntropy(arena: Allocator, inputs: []const Type, attrs: []const Attribute) InferError![]const Type {675    _ = attrs;676    if (inputs.len != 2) return error.ArityMismatch;677    const logits = inputs[0];678    const targets = inputs[1];679    if (!logits.dtype.isFloat()) return error.DTypeMismatch;680    if (!isIndexDType(targets.dtype)) return error.DTypeMismatch;681    if (logits.dims.len != 2 or targets.dims.len != 1) return error.RankMismatch;682    if (logits.dims[0] < 0 or logits.dims[1] < 0) return error.InvalidDimension;683    if (targets.dims[0] != logits.dims[0]) return error.ShapeMismatch;684685    const out_dims = try arena.dupe(i64, logits.dims[0..1]);686    return singletonType(arena, .{ .dtype = logits.dtype, .dims = out_dims });687}688689pub fn inferPad(arena: Allocator, inputs: []const Type, attrs: []const Attribute) InferError![]const Type {690    if (inputs.len != 2) return error.ArityMismatch;691    const low = try inferAttrI64List(attrs, 0);692    const high = try inferAttrI64List(attrs, 1);693    const interior = try inferAttrI64List(attrs, 2);694    const operand = inputs[0];695    const padding_value = inputs[1];696    if (padding_value.dtype != operand.dtype) return error.DTypeMismatch;697    if (padding_value.dims.len != 0) return error.RankMismatch;698    if (low.len != operand.dims.len or high.len != operand.dims.len or interior.len != operand.dims.len) {699        return error.InvalidDimension;700    }701702    const out_dims = try arena.alloc(i64, operand.dims.len);703    for (operand.dims, low, high, interior, 0..) |input_dim, lo, hi, inner, i| {704        if (input_dim < 0 or inner < 0) return error.InvalidDimension;705        const gap_count = if (input_dim > 0) input_dim - 1 else 0;706        const interior_total = std.math.mul(i64, gap_count, inner) catch return error.InvalidDimension;707        const with_low = std.math.add(i64, input_dim, lo) catch return error.InvalidDimension;708        const with_high = std.math.add(i64, with_low, hi) catch return error.InvalidDimension;709        const result_dim = std.math.add(i64, with_high, interior_total) catch return error.InvalidDimension;710        if (result_dim < 0) return error.InvalidDimension;711        out_dims[i] = result_dim;712    }713    return singletonType(arena, .{ .dtype = operand.dtype, .dims = out_dims });714}715716fn isIndexDType(dtype: DType) bool {717    return dtype.isSignedInt() or dtype.isUnsignedInt();718}719720pub fn inferConcatenate(arena: Allocator, inputs: []const Type, attrs: []const Attribute) InferError![]const Type {721    if (inputs.len == 0) return error.ArityMismatch;722    const dim = try inferAttrI64(attrs, 0);723    const first = inputs[0];724    const rank: i64 = @intCast(first.dims.len);725    if (dim < 0 or dim >= rank) return error.InvalidDimension;726    const concat_axis: usize = @intCast(dim);727728    var summed = first.dims[concat_axis];729    for (inputs[1..]) |t| {730        if (t.dtype != first.dtype) return error.DTypeMismatch;731        if (t.dims.len != first.dims.len) return error.RankMismatch;732        for (t.dims, first.dims, 0..) |a, b, i| {733            if (i == concat_axis) continue;734            if (a != b) return error.DimMismatch;735        }736        summed += t.dims[concat_axis];737    }738    const out_dims = try arena.alloc(i64, first.dims.len);739    @memcpy(out_dims, first.dims);740    out_dims[concat_axis] = summed;741    return singletonType(arena, .{ .dtype = first.dtype, .dims = out_dims });742}743744pub fn inferScratch(arena: Allocator, inputs: []const Type, attrs: []const Attribute) InferError![]const Type {745    const words = try inferAttrI64(attrs, 0);746    if (inputs.len != 0) return error.ArityMismatch;747    if (words < 1) return error.InvalidDimension;748    const dims = try arena.alloc(i64, 1);749    dims[0] = words;750    return singletonType(arena, .{ .dtype = .u32, .dims = dims });751}752753pub fn inferCumsum(arena: Allocator, inputs: []const Type, attrs: []const Attribute) InferError![]const Type {754    const axis = try inferAttrI64(attrs, 0);755    if (inputs.len < 1 or inputs.len > 2) return error.ArityMismatch;756    if (axis < 0 or axis >= @as(i64, @intCast(inputs[0].dims.len))) return error.InvalidDimension;757    return singletonType(arena, .{758        .dtype = inputs[0].dtype,759        .dims = try duplicateDims(arena, inputs[0].dims),760    });761}762763pub fn inferIterate(arena: Allocator, inputs: []const Type, attrs: []const Attribute) InferError![]const Type {764    const max_iters = try inferAttrI64(attrs, 0);765    if (max_iters < 1) return error.InvalidDimension;766    if (inputs.len == 0) return error.ArityMismatch;767    try verifySingleElementDomain(inputs);768    const results = try arena.alloc(Type, inputs.len);769    for (results, inputs) |*result, input| {770        result.* = .{ .dtype = input.dtype, .dims = try duplicateDims(arena, input.dims) };771    }772    return results;773}774775pub fn inferIterateYield(arena: Allocator, inputs: []const Type, attrs: []const Attribute) InferError![]const Type {776    _ = attrs;777    if (inputs.len < 2) return error.ArityMismatch;778    if (inputs[0].dtype != .i1) return error.DTypeMismatch;779    if (typeElementCount(inputs[0]) != 1) {780        for (inputs[1..]) |carry| {781            if (typeElementCount(carry) == 1) continue;782            if (!std.mem.eql(i64, carry.dims, inputs[0].dims)) return error.ShapeMismatch;783        }784    }785    try verifySingleElementDomain(inputs[1..]);786    return try arena.alloc(Type, 0);787}788789fn verifySingleElementDomain(inputs: []const Type) InferError!void {790    var domain: ?u64 = null;791    for (inputs) |input| {792        const count = typeElementCount(input);793        if (count == 1) continue;794        if (domain) |existing| {795            if (count != existing) return error.ShapeMismatch;796        } else {797            domain = count;798        }799    }800}801802fn typeElementCount(input: Type) u64 {803    var count: u64 = 1;804    for (input.dims) |dim| {805        if (dim <= 0) return 0;806        count *= @intCast(dim);807    }808    return count;809}810811pub fn inferSelect(arena: Allocator, inputs: []const Type, attrs: []const Attribute) InferError![]const Type {812    _ = attrs;813    if (inputs.len != 3) return error.ArityMismatch;814    const pred = inputs[0];815    const ot = inputs[1];816    const of = inputs[2];817    if (pred.dtype != .i1) return error.DTypeMismatch;818    if (ot.dtype != of.dtype) return error.DTypeMismatch;819    if (!std.mem.eql(i64, ot.dims, of.dims)) return error.ShapeMismatch;820    if (pred.dims.len != 0 and !std.mem.eql(i64, pred.dims, ot.dims)) {821        return error.ShapeMismatch;822    }823    return singletonType(arena, .{824        .dtype = ot.dtype,825        .dims = try duplicateDims(arena, ot.dims),826    });827}828829pub fn inferKernelCall(arena: Allocator, inputs: []const Type, attrs: []const Attribute) InferError![]const Type {830    const contract = try inferAttrKernelCall(attrs, 0);831    if (contract.target.len == 0) return error.InvalidKernelContract;832    if (contract.version == 0) return error.InvalidKernelContract;833    if (contract.results.len == 0) return error.InvalidKernelContract;834    if (contract.operand_effects.len != inputs.len) return error.InvalidKernelContract;835    if (contract.result_aliases.len != contract.results.len) return error.InvalidKernelContract;836    const aliased_operands = try arena.alloc(bool, inputs.len);837    @memset(aliased_operands, false);838    for (contract.result_aliases, 0..) |alias, result_index| {839        const operand_index = alias orelse continue;840        if (operand_index >= inputs.len) return error.InvalidKernelContract;841        if (aliased_operands[operand_index]) return error.InvalidKernelContract;842        aliased_operands[operand_index] = true;843        if (!contract.operand_effects[operand_index].writes()) return error.InvalidKernelContract;844        if (!contract.results[result_index].eql(inputs[operand_index])) return error.InvalidKernelContract;845    }846    return try duplicateTypes(arena, contract.results);847}848849pub fn inferReturn(arena: Allocator, inputs: []const Type, attrs: []const Attribute) InferError![]const Type {850    _ = inputs;851    _ = attrs;852    return try arena.alloc(Type, 0);853}854855const testArena = if (@import("builtin").is_test)856    struct {857        fn init(backing: std.mem.Allocator) std.heap.ArenaAllocator {858            return std.heap.ArenaAllocator.init(backing);859        }860    }.init861else {};862863test "add infers result type from matching operands" {864    var a = testArena(testing.allocator);865    defer a.deinit();866    const dims = [_]i64{ 4, 8 };867    const t: Type = .{ .dtype = .f32, .dims = &dims };868    const out = try inferShape(.add, a.allocator(), &.{ t, t }, &.{});869    try testing.expectEqual(@as(usize, 1), out.len);870    try testing.expect(out[0].eql(t));871}872873test "add rejects dtype mismatch" {874    var a = testArena(testing.allocator);875    defer a.deinit();876    const dims = [_]i64{ 2, 3 };877    const t0: Type = .{ .dtype = .f32, .dims = &dims };878    const t1: Type = .{ .dtype = .i32, .dims = &dims };879    try testing.expectError(error.DTypeMismatch, inferShape(.add, a.allocator(), &.{ t0, t1 }, &.{}));880}881882test "add rejects shape mismatch" {883    var a = testArena(testing.allocator);884    defer a.deinit();885    const d0 = [_]i64{ 2, 3 };886    const d1 = [_]i64{ 2, 4 };887    const t0: Type = .{ .dtype = .f32, .dims = &d0 };888    const t1: Type = .{ .dtype = .f32, .dims = &d1 };889    try testing.expectError(error.ShapeMismatch, inferShape(.add, a.allocator(), &.{ t0, t1 }, &.{}));890}891892test "elementwise binary rejects wrong arity" {893    var a = testArena(testing.allocator);894    defer a.deinit();895    const d = [_]i64{3};896    const t: Type = .{ .dtype = .f32, .dims = &d };897    try testing.expectError(error.ArityMismatch, inferShape(.mul, a.allocator(), &.{t}, &.{}));898}899900test "compare emits i1 with operand shape" {901    var a = testArena(testing.allocator);902    defer a.deinit();903    const dims = [_]i64{ 4, 8 };904    const t: Type = .{ .dtype = .f32, .dims = &dims };905    const attrs = [_]Attribute{.{ .compare_direction = .lt }};906    const out = try inferShape(.compare, a.allocator(), &.{ t, t }, &attrs);907    try testing.expectEqual(@as(usize, 1), out.len);908    try testing.expectEqual(@as(DType, .i1), out[0].dtype);909    try testing.expectEqualSlices(i64, &dims, out[0].dims);910}911912test "compare requires a compare_direction attribute" {913    var a = testArena(testing.allocator);914    defer a.deinit();915    const dims = [_]i64{2};916    const t: Type = .{ .dtype = .f32, .dims = &dims };917    try testing.expectError(error.AttributeMissing, inferShape(.compare, a.allocator(), &.{ t, t }, &.{}));918}919920test "neg preserves type" {921    var a = testArena(testing.allocator);922    defer a.deinit();923    const dims = [_]i64{6};924    const t: Type = .{ .dtype = .f32, .dims = &dims };925    const out = try inferShape(.neg, a.allocator(), &.{t}, &.{});926    try testing.expect(out[0].eql(t));927}928929test "activation preserves type and requires activation kind" {930    var a = testArena(testing.allocator);931    defer a.deinit();932    const dims = [_]i64{8};933    const t: Type = .{ .dtype = .f32, .dims = &dims };934    const attrs = [_]Attribute{.{ .activation_kind = .gelu }};935    const out = try inferShape(.activation, a.allocator(), &.{t}, &attrs);936    try testing.expect(out[0].eql(t));937    try testing.expectError(error.AttributeMissing, inferShape(.activation, a.allocator(), &.{t}, &.{}));938}939940test "convert changes dtype, keeps shape" {941    var a = testArena(testing.allocator);942    defer a.deinit();943    const dims = [_]i64{ 3, 4 };944    const t: Type = .{ .dtype = .i32, .dims = &dims };945    const attrs = [_]Attribute{.{ .dtype = .f32 }};946    const out = try inferShape(.convert, a.allocator(), &.{t}, &attrs);947    try testing.expectEqual(@as(DType, .f32), out[0].dtype);948    try testing.expectEqualSlices(i64, &dims, out[0].dims);949}950951test "constant reads dtype + dims from attributes" {952    var a = testArena(testing.allocator);953    defer a.deinit();954    var payload: [16]u8 = @splat(0);955    const dims = [_]i64{ 2, 2 };956    const attrs = [_]Attribute{957        .{ .bytes = &payload },958        .{ .dtype = .f32 },959        .{ .i64_list = &dims },960    };961    const out = try inferShape(.constant, a.allocator(), &.{}, &attrs);962    try testing.expectEqual(@as(DType, .f32), out[0].dtype);963    try testing.expectEqualSlices(i64, &dims, out[0].dims);964}965966test "iota rejects out-of-range iota_dimension" {967    var a = testArena(testing.allocator);968    defer a.deinit();969    const dims = [_]i64{ 3, 4 };970    const attrs = [_]Attribute{971        .{ .i64 = 2 },972        .{ .dtype = .i32 },973        .{ .i64_list = &dims },974    };975    try testing.expectError(error.InvalidDimension, inferShape(.iota, a.allocator(), &.{}, &attrs));976}977978test "parameter carries dtype + dims" {979    var a = testArena(testing.allocator);980    defer a.deinit();981    const dims = [_]i64{ 4, 8 };982    const attrs = [_]Attribute{983        .{ .i64 = 0 },984        .{ .dtype = .f32 },985        .{ .i64_list = &dims },986    };987    const out = try inferShape(.parameter, a.allocator(), &.{}, &attrs);988    try testing.expectEqualSlices(i64, &dims, out[0].dims);989}990991test "reduce drops reduction dims" {992    var a = testArena(testing.allocator);993    defer a.deinit();994    const in_dims = [_]i64{ 2, 3, 4 };995    const init_dims = [_]i64{};996    const red_dims = [_]i64{1};997    const input: Type = .{ .dtype = .f32, .dims = &in_dims };998    const init_v: Type = .{ .dtype = .f32, .dims = &init_dims };999    const attrs = [_]Attribute{1000        .{ .reducer_kind = .sum },1001        .{ .i64_list = &red_dims },1002    };1003    const out = try inferShape(.reduce, a.allocator(), &.{ input, init_v }, &attrs);1004    try testing.expectEqualSlices(i64, &[_]i64{ 2, 4 }, out[0].dims);1005}10061007test "reduce rejects non-scalar init value" {1008    var a = testArena(testing.allocator);1009    defer a.deinit();1010    const in_dims = [_]i64{ 2, 3 };1011    const init_dims = [_]i64{1};1012    const red_dims = [_]i64{0};1013    const input: Type = .{ .dtype = .f32, .dims = &in_dims };1014    const init_v: Type = .{ .dtype = .f32, .dims = &init_dims };1015    const attrs = [_]Attribute{1016        .{ .reducer_kind = .sum },1017        .{ .i64_list = &red_dims },1018    };1019    try testing.expectError(error.RankMismatch, inferShape(.reduce, a.allocator(), &.{ input, init_v }, &attrs));1020}10211022test "dot_general matmul: [M,K] x [K,N] -> [M,N]" {1023    var a = testArena(testing.allocator);1024    defer a.deinit();1025    const lhs_d = [_]i64{ 4, 8 };1026    const rhs_d = [_]i64{ 8, 16 };1027    const lhs: Type = .{ .dtype = .f32, .dims = &lhs_d };1028    const rhs: Type = .{ .dtype = .f32, .dims = &rhs_d };1029    const empty = [_]i64{};1030    const lhs_c = [_]i64{1};1031    const rhs_c = [_]i64{0};1032    const attrs = [_]Attribute{1033        .{ .i64_list = &empty },1034        .{ .i64_list = &empty },1035        .{ .i64_list = &lhs_c },1036        .{ .i64_list = &rhs_c },1037    };1038    const out = try inferShape(.dot_general, a.allocator(), &.{ lhs, rhs }, &attrs);1039    try testing.expectEqualSlices(i64, &[_]i64{ 4, 16 }, out[0].dims);1040}10411042test "dot_general batched: [B,M,K] x [B,K,N] -> [B,M,N]" {1043    var a = testArena(testing.allocator);1044    defer a.deinit();1045    const lhs_d = [_]i64{ 2, 3, 5 };1046    const rhs_d = [_]i64{ 2, 5, 7 };1047    const lhs: Type = .{ .dtype = .f32, .dims = &lhs_d };1048    const rhs: Type = .{ .dtype = .f32, .dims = &rhs_d };1049    const b_dims = [_]i64{0};1050    const lhs_c = [_]i64{2};1051    const rhs_c = [_]i64{1};1052    const attrs = [_]Attribute{1053        .{ .i64_list = &b_dims },1054        .{ .i64_list = &b_dims },1055        .{ .i64_list = &lhs_c },1056        .{ .i64_list = &rhs_c },1057    };1058    const out = try inferShape(.dot_general, a.allocator(), &.{ lhs, rhs }, &attrs);1059    try testing.expectEqualSlices(i64, &[_]i64{ 2, 3, 7 }, out[0].dims);1060}10611062test "dot_general rejects mismatched contracting dims" {1063    var a = testArena(testing.allocator);1064    defer a.deinit();1065    const lhs_d = [_]i64{ 4, 8 };1066    const rhs_d = [_]i64{ 9, 16 };1067    const lhs: Type = .{ .dtype = .f32, .dims = &lhs_d };1068    const rhs: Type = .{ .dtype = .f32, .dims = &rhs_d };1069    const empty = [_]i64{};1070    const lhs_c = [_]i64{1};1071    const rhs_c = [_]i64{0};1072    const attrs = [_]Attribute{1073        .{ .i64_list = &empty },1074        .{ .i64_list = &empty },1075        .{ .i64_list = &lhs_c },1076        .{ .i64_list = &rhs_c },1077    };1078    try testing.expectError(error.ContractionMismatch, inferShape(.dot_general, a.allocator(), &.{ lhs, rhs }, &attrs));1079}10801081test "einsum infers matrix product output shape" {1082    var a = testArena(testing.allocator);1083    defer a.deinit();1084    const lhs_d = [_]i64{ 4, 8 };1085    const rhs_d = [_]i64{ 8, 16 };1086    const lhs: Type = .{ .dtype = .f32, .dims = &lhs_d };1087    const rhs: Type = .{ .dtype = .f32, .dims = &rhs_d };1088    const attrs = [_]Attribute{.{ .einsum = "ik,kj->ij" }};1089    const out = try inferShape(.einsum, a.allocator(), &.{ lhs, rhs }, &attrs);1090    try testing.expectEqualSlices(i64, &[_]i64{ 4, 16 }, out[0].dims);1091}10921093test "einsum infers scalar reduction output shape" {1094    var a = testArena(testing.allocator);1095    defer a.deinit();1096    const dims = [_]i64{ 4, 8 };1097    const typ: Type = .{ .dtype = .f32, .dims = &dims };1098    const attrs = [_]Attribute{.{ .einsum = "ij->" }};1099    const out = try inferShape(.einsum, a.allocator(), &.{typ}, &attrs);1100    try testing.expectEqual(@as(usize, 0), out[0].dims.len);1101}11021103test "einsum rejects inconsistent shared dimensions" {1104    var a = testArena(testing.allocator);1105    defer a.deinit();1106    const lhs_d = [_]i64{ 4, 8 };1107    const rhs_d = [_]i64{ 9, 16 };1108    const lhs: Type = .{ .dtype = .f32, .dims = &lhs_d };1109    const rhs: Type = .{ .dtype = .f32, .dims = &rhs_d };1110    const attrs = [_]Attribute{.{ .einsum = "ik,kj->ij" }};1111    try testing.expectError(error.InvalidEinsum, inferShape(.einsum, a.allocator(), &.{ lhs, rhs }, &attrs));1112}11131114test "broadcast prepends leading dims" {1115    var a = testArena(testing.allocator);1116    defer a.deinit();1117    const in_d = [_]i64{ 3, 4 };1118    const t: Type = .{ .dtype = .f32, .dims = &in_d };1119    const sizes = [_]i64{2};1120    const attrs = [_]Attribute{.{ .i64_list = &sizes }};1121    const out = try inferShape(.broadcast, a.allocator(), &.{t}, &attrs);1122    try testing.expectEqualSlices(i64, &[_]i64{ 2, 3, 4 }, out[0].dims);1123}11241125test "broadcast_in_dim matches StableHLO 1x3 -> 2x3x2 example" {1126    var a = testArena(testing.allocator);1127    defer a.deinit();1128    const in_d = [_]i64{ 1, 3 };1129    const t: Type = .{ .dtype = .i32, .dims = &in_d };1130    const bcast = [_]i64{ 2, 1 };1131    const result = [_]i64{ 2, 3, 2 };1132    const attrs = [_]Attribute{1133        .{ .i64_list = &bcast },1134        .{ .i64_list = &result },1135    };1136    const out = try inferShape(.broadcast_in_dim, a.allocator(), &.{t}, &attrs);1137    try testing.expectEqualSlices(i64, &result, out[0].dims);1138}11391140test "broadcast_in_dim rejects incompatible source dim" {1141    var a = testArena(testing.allocator);1142    defer a.deinit();1143    const in_d = [_]i64{ 2, 3 };1144    const t: Type = .{ .dtype = .i32, .dims = &in_d };1145    const bcast = [_]i64{ 0, 1 };1146    const result = [_]i64{ 4, 3 };1147    const attrs = [_]Attribute{1148        .{ .i64_list = &bcast },1149        .{ .i64_list = &result },1150    };1151    try testing.expectError(error.ShapeMismatch, inferShape(.broadcast_in_dim, a.allocator(), &.{t}, &attrs));1152}11531154test "reshape changes rank, preserves element count" {1155    var a = testArena(testing.allocator);1156    defer a.deinit();1157    const in_d = [_]i64{ 2, 3 };1158    const t: Type = .{ .dtype = .i32, .dims = &in_d };1159    const new = [_]i64{ 3, 2 };1160    const attrs = [_]Attribute{.{ .i64_list = &new }};1161    const out = try inferShape(.reshape, a.allocator(), &.{t}, &attrs);1162    try testing.expectEqualSlices(i64, &new, out[0].dims);1163}11641165test "reshape rejects element-count mismatch" {1166    var a = testArena(testing.allocator);1167    defer a.deinit();1168    const in_d = [_]i64{ 2, 3 };1169    const t: Type = .{ .dtype = .i32, .dims = &in_d };1170    const new = [_]i64{ 4, 2 };1171    const attrs = [_]Attribute{.{ .i64_list = &new }};1172    try testing.expectError(error.ElementCountMismatch, inferShape(.reshape, a.allocator(), &.{t}, &attrs));1173}11741175test "transpose permutes dims" {1176    var a = testArena(testing.allocator);1177    defer a.deinit();1178    const in_d = [_]i64{ 2, 3, 4 };1179    const t: Type = .{ .dtype = .f32, .dims = &in_d };1180    const perm = [_]i64{ 2, 0, 1 };1181    const attrs = [_]Attribute{.{ .i64_list = &perm }};1182    const out = try inferShape(.transpose, a.allocator(), &.{t}, &attrs);1183    try testing.expectEqualSlices(i64, &[_]i64{ 4, 2, 3 }, out[0].dims);1184}11851186test "transpose rejects non-permutation" {1187    var a = testArena(testing.allocator);1188    defer a.deinit();1189    const in_d = [_]i64{ 2, 3, 4 };1190    const t: Type = .{ .dtype = .f32, .dims = &in_d };1191    const perm = [_]i64{ 0, 0, 1 };1192    const attrs = [_]Attribute{.{ .i64_list = &perm }};1193    try testing.expectError(error.InvalidPermutation, inferShape(.transpose, a.allocator(), &.{t}, &attrs));1194}11951196test "slice computes ceil((limit - start) / stride)" {1197    var a = testArena(testing.allocator);1198    defer a.deinit();1199    const in_d = [_]i64{ 10, 10 };1200    const t: Type = .{ .dtype = .i32, .dims = &in_d };1201    const start = [_]i64{ 1, 2 };1202    const limit = [_]i64{ 9, 8 };1203    const stride = [_]i64{ 2, 3 };1204    const attrs = [_]Attribute{1205        .{ .i64_list = &start },1206        .{ .i64_list = &limit },1207        .{ .i64_list = &stride },1208    };1209    const out = try inferShape(.slice, a.allocator(), &.{t}, &attrs);1210    try testing.expectEqualSlices(i64, &[_]i64{ 4, 2 }, out[0].dims);1211}12121213test "gather inserts index shape at the selected axis" {1214    var a = testArena(testing.allocator);1215    defer a.deinit();1216    const operand_dims = [_]i64{ 2, 3, 4 };1217    const index_dims = [_]i64{ 5, 6 };1218    const operand: Type = .{ .dtype = .f32, .dims = &operand_dims };1219    const indices: Type = .{ .dtype = .i32, .dims = &index_dims };1220    const attrs = [_]Attribute{.{ .i64 = 1 }};1221    const out = try inferShape(.gather, a.allocator(), &.{ operand, indices }, &attrs);1222    try testing.expectEqualSlices(i64, &[_]i64{ 2, 5, 6, 4 }, out[0].dims);1223}12241225test "gather rejects invalid axis and non-integer indices" {1226    var a = testArena(testing.allocator);1227    defer a.deinit();1228    const operand_dims = [_]i64{ 2, 3 };1229    const index_dims = [_]i64{4};1230    const operand: Type = .{ .dtype = .f32, .dims = &operand_dims };1231    const indices: Type = .{ .dtype = .i32, .dims = &index_dims };1232    const bad_indices: Type = .{ .dtype = .f32, .dims = &index_dims };1233    try testing.expectError(error.InvalidDimension, inferShape(.gather, a.allocator(), &.{ operand, indices }, &.{.{ .i64 = 2 }}));1234    try testing.expectError(error.DTypeMismatch, inferShape(.gather, a.allocator(), &.{ operand, bad_indices }, &.{.{ .i64 = 0 }}));1235}12361237test "scatter returns operand shape and checks update shape" {1238    var a = testArena(testing.allocator);1239    defer a.deinit();1240    const operand_dims = [_]i64{ 2, 3, 4 };1241    const index_dims = [_]i64{ 5, 6 };1242    const update_dims = [_]i64{ 2, 5, 6, 4 };1243    const bad_update_dims = [_]i64{ 2, 5, 4 };1244    const operand: Type = .{ .dtype = .f32, .dims = &operand_dims };1245    const indices: Type = .{ .dtype = .i32, .dims = &index_dims };1246    const updates: Type = .{ .dtype = .f32, .dims = &update_dims };1247    const bad_updates: Type = .{ .dtype = .f32, .dims = &bad_update_dims };1248    const attrs = [_]Attribute{.{ .i64 = 1 }};1249    const out = try inferShape(.scatter, a.allocator(), &.{ operand, indices, updates }, &attrs);1250    try testing.expectEqualSlices(i64, &operand_dims, out[0].dims);1251    try testing.expectError(error.RankMismatch, inferShape(.scatter, a.allocator(), &.{ operand, indices, bad_updates }, &attrs));1252}12531254test "scatter rejects invalid axis and non-integer indices" {1255    var a = testArena(testing.allocator);1256    defer a.deinit();1257    const operand_dims = [_]i64{ 2, 3 };1258    const index_dims = [_]i64{4};1259    const update_dims = [_]i64{ 4, 3 };1260    const operand: Type = .{ .dtype = .i32, .dims = &operand_dims };1261    const indices: Type = .{ .dtype = .i32, .dims = &index_dims };1262    const bad_indices: Type = .{ .dtype = .f32, .dims = &index_dims };1263    const updates: Type = .{ .dtype = .i32, .dims = &update_dims };1264    try testing.expectError(error.InvalidDimension, inferShape(.scatter, a.allocator(), &.{ operand, indices, updates }, &.{.{ .i64 = 2 }}));1265    try testing.expectError(error.DTypeMismatch, inferShape(.scatter, a.allocator(), &.{ operand, bad_indices, updates }, &.{.{ .i64 = 0 }}));1266}12671268test "pad infers positive edge and interior padding" {1269    var a = testArena(testing.allocator);1270    defer a.deinit();1271    const in_d = [_]i64{ 2, 3 };1272    const scalar_d = [_]i64{};1273    const input: Type = .{ .dtype = .i32, .dims = &in_d };1274    const padding_value: Type = .{ .dtype = .i32, .dims = &scalar_d };1275    const low = [_]i64{ 1, 0 };1276    const high = [_]i64{ 0, 2 };1277    const interior = [_]i64{ 0, 1 };1278    const attrs = [_]Attribute{1279        .{ .i64_list = &low },1280        .{ .i64_list = &high },1281        .{ .i64_list = &interior },1282    };1283    const out = try inferShape(.pad, a.allocator(), &.{ input, padding_value }, &attrs);1284    try testing.expectEqualSlices(i64, &[_]i64{ 3, 7 }, out[0].dims);1285}12861287test "pad allows negative edge padding when result shape remains valid" {1288    var a = testArena(testing.allocator);1289    defer a.deinit();1290    const in_d = [_]i64{4};1291    const scalar_d = [_]i64{};1292    const input: Type = .{ .dtype = .f32, .dims = &in_d };1293    const padding_value: Type = .{ .dtype = .f32, .dims = &scalar_d };1294    const low = [_]i64{-1};1295    const high = [_]i64{2};1296    const interior = [_]i64{0};1297    const attrs = [_]Attribute{1298        .{ .i64_list = &low },1299        .{ .i64_list = &high },1300        .{ .i64_list = &interior },1301    };1302    const out = try inferShape(.pad, a.allocator(), &.{ input, padding_value }, &attrs);1303    try testing.expectEqualSlices(i64, &[_]i64{5}, out[0].dims);1304}13051306test "pad rejects non-scalar padding value and invalid interior padding" {1307    var a = testArena(testing.allocator);1308    defer a.deinit();1309    const in_d = [_]i64{ 2, 3 };1310    const scalar_d = [_]i64{};1311    const vector_d = [_]i64{1};1312    const input: Type = .{ .dtype = .i32, .dims = &in_d };1313    const scalar_padding: Type = .{ .dtype = .i32, .dims = &scalar_d };1314    const vector_padding: Type = .{ .dtype = .i32, .dims = &vector_d };1315    const low = [_]i64{ 0, 0 };1316    const high = [_]i64{ 0, 0 };1317    const interior = [_]i64{ 0, -1 };1318    const attrs = [_]Attribute{1319        .{ .i64_list = &low },1320        .{ .i64_list = &high },1321        .{ .i64_list = &low },1322    };1323    try testing.expectError(error.RankMismatch, inferShape(.pad, a.allocator(), &.{ input, vector_padding }, &attrs));13241325    const bad_attrs = [_]Attribute{1326        .{ .i64_list = &low },1327        .{ .i64_list = &high },1328        .{ .i64_list = &interior },1329    };1330    try testing.expectError(error.InvalidDimension, inferShape(.pad, a.allocator(), &.{ input, scalar_padding }, &bad_attrs));1331}13321333test "concatenate sums along the concat dim" {1334    var a = testArena(testing.allocator);1335    defer a.deinit();1336    const d0 = [_]i64{ 3, 2 };1337    const d1 = [_]i64{ 1, 2 };1338    const t0: Type = .{ .dtype = .i64, .dims = &d0 };1339    const t1: Type = .{ .dtype = .i64, .dims = &d1 };1340    const attrs = [_]Attribute{.{ .i64 = 0 }};1341    const out = try inferShape(.concatenate, a.allocator(), &.{ t0, t1 }, &attrs);1342    try testing.expectEqualSlices(i64, &[_]i64{ 4, 2 }, out[0].dims);1343}13441345test "concatenate rejects non-concat-dim mismatch" {1346    var a = testArena(testing.allocator);1347    defer a.deinit();1348    const d0 = [_]i64{ 3, 2 };1349    const d1 = [_]i64{ 1, 3 };1350    const t0: Type = .{ .dtype = .i64, .dims = &d0 };1351    const t1: Type = .{ .dtype = .i64, .dims = &d1 };1352    const attrs = [_]Attribute{.{ .i64 = 0 }};1353    try testing.expectError(error.DimMismatch, inferShape(.concatenate, a.allocator(), &.{ t0, t1 }, &attrs));1354}13551356test "select requires i1 predicate and matching branches" {1357    var a = testArena(testing.allocator);1358    defer a.deinit();1359    const d = [_]i64{ 2, 2 };1360    const pred: Type = .{ .dtype = .i1, .dims = &d };1361    const on_true: Type = .{ .dtype = .i32, .dims = &d };1362    const on_false: Type = .{ .dtype = .i32, .dims = &d };1363    const out = try inferShape(.select, a.allocator(), &.{ pred, on_true, on_false }, &.{});1364    try testing.expect(out[0].eql(on_true));13651366    const bad_pred: Type = .{ .dtype = .i32, .dims = &d };1367    try testing.expectError(error.DTypeMismatch, inferShape(.select, a.allocator(), &.{ bad_pred, on_true, on_false }, &.{}));1368}13691370test "kernel_call returns explicit contract result types" {1371    var a = testArena(testing.allocator);1372    defer a.deinit();1373    const input_dims = [_]i64{ 4, 8 };1374    const result_dims = [_]i64{ 4, 8 };1375    const input: Type = .{ .dtype = .f32, .dims = &input_dims };1376    const result: Type = .{ .dtype = .f32, .dims = &result_dims };1377    const contract = KernelCallContract{1378        .target = "scale_f32",1379        .version = 1,1380        .has_side_effects = false,1381        .operand_effects = &.{.none},1382        .result_aliases = &.{null},1383        .results = &.{result},1384    };1385    const out = try inferShape(.kernel_call, a.allocator(), &.{input}, &.{.{ .kernel_call = contract }});1386    try testing.expectEqual(@as(usize, 1), out.len);1387    try testing.expect(out[0].eql(result));1388}13891390test "kernel_call accepts explicit result aliases for writable operands" {1391    var a = testArena(testing.allocator);1392    defer a.deinit();1393    const dims = [_]i64{4};1394    const input: Type = .{ .dtype = .f32, .dims = &dims };1395    const contract = KernelCallContract{1396        .target = "update_f32",1397        .version = 1,1398        .has_side_effects = false,1399        .operand_effects = &.{.read_write},1400        .result_aliases = &.{0},1401        .results = &.{input},1402    };1403    const out = try inferShape(.kernel_call, a.allocator(), &.{input}, &.{.{ .kernel_call = contract }});1404    try testing.expectEqual(@as(usize, 1), out.len);1405    try testing.expect(out[0].eql(input));1406}14071408test "kernel_call rejects malformed contracts" {1409    var a = testArena(testing.allocator);1410    defer a.deinit();1411    const dims = [_]i64{4};1412    const input: Type = .{ .dtype = .f32, .dims = &dims };1413    const result: Type = .{ .dtype = .f32, .dims = &dims };1414    try testing.expectError(1415        error.InvalidKernelContract,1416        inferShape(.kernel_call, a.allocator(), &.{}, &.{.{ .kernel_call = .{1417            .target = "",1418            .version = 1,1419            .has_side_effects = false,1420            .operand_effects = &.{},1421            .result_aliases = &.{null},1422            .results = &.{result},1423        } }}),1424    );1425    try testing.expectError(1426        error.InvalidKernelContract,1427        inferShape(.kernel_call, a.allocator(), &.{}, &.{.{ .kernel_call = .{1428            .target = "scale_f32",1429            .version = 0,1430            .has_side_effects = false,1431            .operand_effects = &.{},1432            .result_aliases = &.{null},1433            .results = &.{result},1434        } }}),1435    );1436    try testing.expectError(1437        error.InvalidKernelContract,1438        inferShape(.kernel_call, a.allocator(), &.{input}, &.{.{ .kernel_call = .{1439            .target = "scale_f32",1440            .version = 1,1441            .has_side_effects = false,1442            .operand_effects = &.{},1443            .result_aliases = &.{null},1444            .results = &.{result},1445        } }}),1446    );1447    try testing.expectError(1448        error.InvalidKernelContract,1449        inferShape(.kernel_call, a.allocator(), &.{input}, &.{.{ .kernel_call = .{1450            .target = "update_f32",1451            .version = 1,1452            .has_side_effects = false,1453            .operand_effects = &.{.read},1454            .result_aliases = &.{0},1455            .results = &.{result},1456        } }}),1457    );1458    try testing.expectError(1459        error.InvalidKernelContract,1460        inferShape(.kernel_call, a.allocator(), &.{input}, &.{.{ .kernel_call = .{1461            .target = "update_f32",1462            .version = 1,1463            .has_side_effects = false,1464            .operand_effects = &.{.read_write},1465            .result_aliases = &.{ 0, 0 },1466            .results = &.{ result, result },1467        } }}),1468    );1469}14701471test "return has zero outputs" {1472    var a = testArena(testing.allocator);1473    defer a.deinit();1474    const d = [_]i64{3};1475    const t: Type = .{ .dtype = .f32, .dims = &d };1476    const out = try inferShape(.@"return", a.allocator(), &.{t}, &.{});1477    try testing.expectEqual(@as(usize, 0), out.len);1478}