def plan_compatible

bendtorch-core@0.1.0.0 · src/tensor/matmul.bend:39

import Base
import bendtorch-core@0.1.0.0/src/shape/core.bend as Shape
import bendtorch-core@0.1.0.0/src/tensor/lifecycle.bend as Lifecycle

39def plan_compatible(-left: Shape.Shape, -right: Shape.Shape, +rows: U32, +inner: U32,
40  +right_inner: U32, +columns: U32,
41  -le: {Shape.dimensions(left) == [rows, inner] : List<&2, U32>},
42  -re: {Shape.dimensions(right) == [right_inner, columns] : List<&2, U32>},
43  decision: Shape.BoolCase<Nat.is_eq(U32.to_nat(inner), U32.to_nat(right_inner))>) ->
44  Result<&2, &2, Error, Plan<left, right>>:
45  match decision:
46    case Shape.BFalse{equal}:
47      Fail{InnerDimensionMismatch{}}
48    case Shape.BTrue{equal}:
49      plan_output(left, right, rows, inner, right_inner, columns, le, re,
50        Lifecycle.natural_equal(U32.to_nat(inner), U32.to_nat(right_inner), equal),
51        Shape.construct([rows, columns]))