def geometry_shaped

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

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

729def geometry_shaped(+lr: Nat, +l1: U32, +l2: U32, +rows: U32, +inner: U32, +lv: Bool, +rr: Nat, +r1: U32, +r2: U32, +columns: U32,
730  +rinner: U32, +rv: Bool, +o1: U32, +o2: U32, +same: {U32.to_nat(inner) == U32.to_nat(rinner) : Nat},
731  +e1: {U32.to_nat(o1) == Laws.extent(U32.to_nat(l1), U32.to_nat(r1)) : Nat},
732  +e2: {U32.to_nat(o2) == Laws.extent(U32.to_nat(l2), U32.to_nat(r2)) : Nat},
733  result: Shape.ConstructResult<List.append(&2, U32, Batched.batch_dims(Batched.larger(lr, rr), o1, o2), Batched.matrix_dims(rows, lv, columns, rv))>) ->
734  ExpectedOf(lr, l1, l2, rows, inner, lv, rr, r1, r2, columns, rinner, rv,
735    Batched.planned_shape(Batched.batch_dims(Batched.larger(lr, rr), o1, o2), rows, inner, columns, l1, l2, r1, r2, o1, o2, lv, rv,
736      List.append(&2, U32, Batched.batch_dims(Batched.larger(lr, rr), o1, o2), Batched.matrix_dims(rows, lv, columns, rv)), result)):
737  match result:
738    case Shape.SError{error}:
739      Unit{}
740    case Shape.SBuilt{+shape, +preserved}:
741      Laws.Geometry{
742        Equal.trans(List<&2, Nat>, Shape.naturals(Shape.dimensions(shape)),
743          Shape.naturals(List.append(&2, U32, Batched.batch_dims(Batched.larger(lr, rr), o1, o2), Batched.matrix_dims(rows, lv, columns, rv))),
744          List.append(&2, Nat, batch_nats(Batched.larger(lr, rr), Laws.extent(U32.to_nat(l1), U32.to_nat(r1)), Laws.extent(U32.to_nat(l2), U32.to_nat(r2))),
745            List.append(&2, Nat, Laws.kept(lv, U32.to_nat(rows)), Laws.kept(rv, U32.to_nat(columns)))),
746          Equal.cong(List<&2, U32>, List<&2, Nat>, dims => Shape.naturals(dims), Shape.dimensions(shape),
747            List.append(&2, U32, Batched.batch_dims(Batched.larger(lr, rr), o1, o2), Batched.matrix_dims(rows, lv, columns, rv)), preserved),
748          output_nats(Batched.larger(lr, rr), o1, o2, Laws.extent(U32.to_nat(l1), U32.to_nat(r1)), Laws.extent(U32.to_nat(l2), U32.to_nat(r2)),
749            e1, e2, rows, lv, columns, rv)),
750        {==},
751        Equal.cong(Nat, List<&2, Nat>, x => [U32.to_nat(r1), U32.to_nat(r2), x, U32.to_nat(columns)], U32.to_nat(inner), U32.to_nat(rinner), same),
752        pair_nats(U32.to_nat(o1), U32.to_nat(o2), Laws.extent(U32.to_nat(l1), U32.to_nat(r1)), Laws.extent(U32.to_nat(l2), U32.to_nat(r2)), e1, e2),
753        {==}}