def geometry_batched

bendtorch-core@0.1.0.0 · src/tensor/batched_matmul_PROOF.bend:755

import Base
import bendtorch-core@0.1.0.0/src/shape/core.bend as Shape
import bendtorch-core@0.1.0.0/src/tensor/matmul.bend as Matrix
import bendtorch-core@0.1.0.0/src/tensor/batched_matmul.bend as Batched

755def geometry_batched(+lr: Nat, +l1: U32, +l2: U32, +rows: U32, +inner: U32, +lv: Bool, +rr: Nat, +r1: U32, +r2: U32, +columns: U32,
756  +rinner: U32, +rv: Bool, +same: {U32.to_nat(inner) == U32.to_nat(rinner) : Nat},
757  first: Result<&2, &2, Matrix.Error, U32>, second: Result<&2, &2, Matrix.Error, U32>,
758  +m1: AxisMeaning(l1, r1, first), +m2: AxisMeaning(l2, r2, second)) ->
759  ExpectedOf(lr, l1, l2, rows, inner, lv, rr, r1, r2, columns, rinner, rv,
760    Batched.batched(Batched.larger(lr, rr), rows, inner, columns, l1, l2, r1, r2, lv, rv, first, second)):
761  match first second:
762    case Fail{error} second:
763      Unit{}
764    case Done{o1} Fail{error}:
765      Unit{}
766    case Done{+o1} Done{+o2}:
767      geometry_shaped(lr, l1, l2, rows, inner, lv, rr, r1, r2, columns, rinner, rv, o1, o2, same, m1, m2,
768        Shape.construct(List.append(&2, U32, Batched.batch_dims(Batched.larger(lr, rr), o1, o2), Batched.matrix_dims(rows, lv, columns, rv))))