def plan_right

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

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

53def plan_right(-left: Shape.Shape, +right: Shape.Shape, +rows: U32, +inner: U32,
54  -le: {Shape.dimensions(left) == [rows, inner] : List<&2, U32>}, dimensions: List<&2, U32>,
55  -same: {dimensions == Shape.dimensions(right) : List<&2, U32>}) ->
56  Result<&2, &2, Error, Plan<left, right>>:
57  match dimensions:
58    case +right_inner <> +columns <> Nil{}:
59      plan_compatible(left, right, rows, inner, right_inner, columns, le,
60        Equal.sym(List<&2, U32>, [right_inner, columns], Shape.dimensions(right), same),
61        Shape.bool_case(Nat.is_eq(U32.to_nat(inner), U32.to_nat(right_inner))))
62    case rest:
63      Fail{RankMismatch{}}