def permuted_count

bendtorch-core@0.1.0.0 · src/tensor/permute_values.bend:339

import Base
import bendtorch-core@0.1.0.0/src/shape/core.bend as Shape
import bendtorch-core@0.1.0.0/src/tensor/lifecycle.bend as Lifecycle
import bendtorch-core@0.1.0.0/src/tensor/layout.bend as Layout
import bendtorch-core@0.1.0.0/src/tensor/from_values_contract.bend as FromValues
import bendtorch-core@0.1.0.0/src/tensor/contracts/replayed.bend as Replayed
import bendtorch-core@0.1.0.0/src/tensor/transform_observation.bend as Observation
import bendtorch-core@0.1.0.0/src/tensor/permute_build_proof.bend as PermuteProof

339def permuted_count(+input: Shape.Shape, +values: List<&2, F32>, +axes: List<&2, U32>, +valid: {True{} == Layout.bijective(input, axes) : Bool},
340  -fits: {Lifecycle.logical_length(Replayed.storage(input, 0.0, FromValues.writes(values, 0n))) == Shape.logical_numel(Shape.dimensions(input)) : Nat}) ->
341  Permuted(input, values, axes, fits, pair => {Observation.applied_count_pair(input, pair) == Some{Shape.count(input)} : Maybe<&2, U32>}):
342  permute_lift(input, values, axes, valid, pair => {Observation.applied_count_pair(input, pair) == Some{Shape.count(input)} : Maybe<&2, U32>},
343    plan => source_fits => target_fits =>
344      Equal.cong(Nat, Maybe<&2, U32>, number => Some{U32.from_nat(number)}, Shape.elements(Layout.output(input, axes, plan)), Shape.elements(input),
345        PermuteProof.plan_elements(input, axes, plan)), fits)