def words_moved

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

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/layout.bend as Layout
import bendtorch-core@0.1.0.0/src/tensor/word_conversion.bend as Conversion
import bendtorch-core@0.1.0.0/src/tensor/contracts/replayed.bend as Replayed
import bendtorch-core@0.1.0.0/src/tensor/contracts/pointwise_proof.bend as PointwiseProof
import bendtorch-core@0.1.0.0/src/tensor/index_map.bend as IndexMap
import bendtorch-core@0.1.0.0/src/tensor/observation.bend as Values
import bendtorch-core@0.1.0.0/src/tensor/permute_values_records.bend as Records
import bendtorch-core@0.1.0.0/src/tensor/contracts/words_proof.bend as Words
import bendtorch-core@0.1.0.0/src/tensor/contracts/LAWS.bend as Contracts

507# The words the copy loop writes, in index order, are `permuted_words`.
508def words_moved(+input: Shape.Shape, +fill: F32, +ops: List<&2, StorageLaws.Op>, +axes: List<&2, U32>, +plan: Layout.Plan<input, axes>,
509  +count: Nat, +room: {Nat.is_le(count, Shape.logical_numel(Shape.dimensions(Layout.output(input, axes, plan)))) == True{} : Bool}) ->
510  {Contracts.words(Layout.output(input, axes, plan), 0.0,
511    IndexMap.writes(~Request, ~source_of, Shape.logical_numel(Shape.dimensions(Layout.output(input, axes, plan))), Request{input, axes}, input, fill, ops),
512    count) == Records.permuted_words(input, fill, ops, axes, count) : List<&2, F32>}:
513  match count:
514    case 0n:
515      {==}
516    case 1n+previous:
517      +output = Layout.output(input, axes, plan)
518      +total = Shape.logical_numel(Shape.dimensions(output))
519      +request : Request = Request{input, axes}
520      +history = IndexMap.writes(~Request, ~source_of, total, request, input, fill, ops)
521      +below = Conversion.le_successor_lt(previous, total, room)
522      +position = U32.from_nat(previous)
523      Words.snoc_equal(Contracts.words(output, 0.0, history, previous), Records.permuted_words(input, fill, ops, axes, previous),
524        Replayed.value(output, 0.0, history, position),
525        Replayed.value(input, fill, ops, U32.from_nat(Values.permute_source_index(input, axes, previous))),
526        words_moved(input, fill, ops, axes, plan, previous, Conversion.le_predecessor(previous, total, room)),
527        Equal.trans(F32, Replayed.value(output, 0.0, history, position),
528          Replayed.value(input, fill, ops, U32.from_nat(Values.permute_source_index(input, axes, U32.to_nat(position)))),
529          Replayed.value(input, fill, ops, U32.from_nat(Values.permute_source_index(input, axes, previous))),
530          PointwiseProof.writes_value(~Unit, ~IndexMap.keep, ~Request, ~source_of, ~source_of, total, Unit{}, request, input, fill, ops,
531            input, fill, ops, output, position, PointwiseProof.word_inside(previous, output, below), Conversion.le_reflexive(total)),
532          Equal.cong(Nat, F32, number => Replayed.value(input, fill, ops, U32.from_nat(Values.permute_source_index(input, axes, number))),
533            U32.to_nat(position), previous, PointwiseProof.word_exact(previous, output, below))))