def batch_part

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

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

664def batch_part(+rank: Nat, +o1: U32, +o2: U32, +first: Nat, +second: Nat,
665  +e1: {U32.to_nat(o1) == first : Nat}, +e2: {U32.to_nat(o2) == second : Nat}) ->
666  {Shape.naturals(Batched.batch_dims(rank, o1, o2)) == batch_nats(rank, first, second) : List<&2, Nat>}:
667  match rank:
668    case 0n:
669      {==}
670    case 1n+rest:
671      match rest:
672        case 0n:
673          %Equal.sym(Nat, U32.to_nat(o2), second, e2) : {[_] == [second] : List<&2, Nat>}
674          {==}
675        case 1n+more:
676          %Equal.sym(Nat, U32.to_nat(o1), first, e1) : {[_, U32.to_nat(o2)] == [first, second] : List<&2, Nat>}
677          %Equal.sym(Nat, U32.to_nat(o2), second, e2) : {[first, _] == [first, second] : List<&2, Nat>}
678          {==}