def zero_head

bendtorch-core@0.1.0.0 · src/tensor/slice_PROOF.bend:375

import Base
import bendtorch-core@0.1.0.0/src/shape/core.bend as Shape
import bendtorch-core@0.1.0.0/src/shape/nat_proofs.bend as Nats
import bendtorch-core@0.1.0.0/src/storage/arith.bend as Arith
import bendtorch-core@0.1.0.0/src/tensor/slice.bend as Slice

375# Some dimension 0 makes the element count 0.
376def zero_head(+extent: U32, +rest: List<&2, U32>, h: {Slice.holds_zero(extent <> rest) == True{} : Bool},
377  rec: @h: {Slice.holds_zero(rest) == True{} : Bool} -> {Shape.logical_numel(rest) == 0n : Nat},
378  decision: Shape.BoolCase<Nat.is_eq(U32.to_nat(extent), 0n)>) -> {Shape.logical_numel(extent <> rest) == 0n : Nat}:
379  match decision:
380    case Shape.BTrue{equal}:
381      %Equal.sym(Nat, U32.to_nat(extent), 0n, Arith.eq_true(U32.to_nat(extent), 0n, Equal.sym(Bool, True{}, Nat.is_eq(U32.to_nat(extent), 0n), equal))) :
382        {Nat.mul(_, Shape.logical_numel(rest)) == 0n : Nat}
383      {==}
384    case Shape.BFalse{equal}:
385      %Equal.sym(Nat, Shape.logical_numel(rest), 0n,
386        rec(Equal.trans(Bool, Slice.holds_zero(rest), Bool.or(Nat.is_eq(U32.to_nat(extent), 0n), Slice.holds_zero(rest)), True{},
387          Equal.cong(Bool, Bool, x => Bool.or(x, Slice.holds_zero(rest)), False{}, Nat.is_eq(U32.to_nat(extent), 0n), equal), h))) :
388        {Nat.mul(U32.to_nat(extent), _) == 0n : Nat}
389      Nats.mul_zero_right(U32.to_nat(extent))