def permuted_kept

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

import Base
import bendtorch-core@0.1.0.0/src/shape/core.bend as Shape
import bendtorch-core@0.1.0.0/src/storage/LAWS.bend as StorageLaws
import bendtorch-core@0.1.0.0/src/tensor/core.bend as Core
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/index_map.bend as IndexMap
import bendtorch-core@0.1.0.0/src/tensor/transform_observation.bend as Observation

369def permuted_kept(+input: Shape.Shape, +values: List<&2, F32>,
370  +counted: {List.length(&2, F32, values) == Shape.logical_numel(Shape.dimensions(input)) : Nat},
371  +axes: List<&2, U32>, +valid: {True{} == Layout.bijective(input, axes) : Bool}, +other: U32,
372  +inside: {Nat.is_lt(U32.to_nat(other), Shape.logical_numel(Shape.dimensions(input))) == True{} : Bool},
373  -fits: {Lifecycle.logical_length(Replayed.storage(input, 0.0, FromValues.writes(values, 0n))) == Shape.logical_numel(Shape.dimensions(input)) : Nat}) ->
374  Permuted(input, values, axes, fits, pair => {Observation.kept_value_pair(input, other, pair) ==
375    List.get(&2, F32, values, U32.to_nat(other)) : Maybe<&2, F32>}):
376  permute_lift(input, values, axes, valid, pair => {Observation.kept_value_pair(input, other, pair) ==
377      List.get(&2, F32, values, U32.to_nat(other)) : Maybe<&2, F32>},
378    plan => source_fits => target_fits =>
379      +total = Shape.logical_numel(Shape.dimensions(Layout.output(input, axes, plan)))
380      +history = FromValues.writes(values, 0n)
381      Equal.trans(Maybe<&2, F32>,
382        Core.read_value(input, Core.read_flat(input, Core.TTensor{kept_source(input, 0.0, FromValues.writes(values, 0n)), source_fits}, other)),
383        Some{Replayed.value(input, 0.0, history, other)}, List.get(&2, F32, values, U32.to_nat(other)),
384        Equal.trans(Maybe<&2, F32>, Core.read_value(input, Core.read_flat(input, Core.TTensor{kept_source(input, 0.0, history), source_fits}, other)),
385          Some{Replayed.value(input, 0.0, Replayed.snoc(history, StorageLaws.First{}), other)}, Some{Replayed.value(input, 0.0, history, other)},
386          Replayed.read_found(input, 0.0, Replayed.snoc(history, StorageLaws.First{}), source_fits, other, inside),
387          Equal.cong(F32, Maybe<&2, F32>, word => Some{word}, Replayed.value(input, 0.0, Replayed.snoc(history, StorageLaws.First{}), other),
388            Replayed.value(input, 0.0, history, other), half_keeps(input, 0.0, history, other))),
389        IndexMap.stored_at(input, values, counted, other, inside)), fits)