def plan_numel

bendtorch-core@0.1.0.0 · src/tensor/permute_build_proof.bend:327

import Base
import bendtorch-core@0.1.0.0/src/shape/core.bend as Shape
import bendtorch-core@0.1.0.0/src/tensor/layout.bend as Layout

327def plan_numel(+input: Shape.Shape, +axes: List<&2, U32>, +plan: Layout.Plan<input, axes>) ->
328  {Shape.logical_numel(Shape.dimensions(Layout.output(input, axes, plan))) == Shape.logical_numel(Shape.dimensions(input)) : Nat}:
329  +output = Layout.output(input, axes, plan)
330  Equal.trans(Nat, Shape.logical_numel(Shape.dimensions(output)), Shape.elements(output), Shape.logical_numel(Shape.dimensions(input)),
331    Equal.sym(Nat, Shape.elements(output), Shape.logical_numel(Shape.dimensions(output)), Shape.elements_exact(output)),
332    Equal.trans(Nat, Shape.elements(output), Shape.elements(input), Shape.logical_numel(Shape.dimensions(input)),
333      plan_elements(input, axes, plan), Shape.elements_exact(input)))