def replayed_final

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

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/layout.bend as Layout
import bendtorch-core@0.1.0.0/src/tensor/transforms.bend as Transforms
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/index_map.bend as IndexMap
import bendtorch-core@0.1.0.0/src/tensor/permute_values_records.bend as Records

535# The copy loop on any replayed source, in transport form: once
536# `Layout.permutation` gives the plan, P holds for the permute if it holds for
537# every result of the form the loop's final state has.
538def replayed_final(+input: Shape.Shape, +fill: F32, +ops: List<&2, StorageLaws.Op>, +axes: List<&2, U32>, +plan: Layout.Plan<input, axes>,
539  -P: Core.Dynamic & Result<&2, &1, Transforms.Error, Core.Dynamic> -> Type,
540  h: Records.PermutedWith(input, fill, ops, axes, plan, P)) ->
541  Final(input, axes, plan, fill, ops, Shape.logical_numel(Shape.dimensions(Layout.output(input, axes, plan))),
542    pair => P(Transforms.dynamic_result(input, Transforms.copy_pack(input, axes, plan, pair)))):
543  +output = Layout.output(input, axes, plan)
544  +total = Shape.logical_numel(Shape.dimensions(output))
545  +request : Request = Request{input, axes}
546  source_fits => target_fits =>
547    h(Replayed.snoc(ops, StorageLaws.First{}), words_first(input, fill, ops, Shape.logical_numel(Shape.dimensions(input))),
548      source_fits,
549      IndexMap.writes(~Request, ~source_of, total, request, input, fill, ops),
550      words_moved(input, fill, ops, axes, plan, total, Conversion.le_reflexive(total)),
551      target_fits)