def numel_pad

bendtorch-core@0.1.0.0 · src/autograd/reduce_parallel_proof.bend:28

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

28def numel_pad(+count: Nat, +dims: List<&2, Nat>) -> {Shape.numel(Shape.pad(count, dims)) == Shape.numel(dims) : Nat}:
29  match count:
30    case 0n:
31      {==}
32    case 1n+rest:
33      Equal.trans(Nat, Nat.add(Shape.numel(Shape.pad(rest, dims)), 0n), Shape.numel(Shape.pad(rest, dims)), Shape.numel(dims),
34        Broadcast.add_zero(Shape.numel(Shape.pad(rest, dims))), numel_pad(rest, dims))