def permuted_at

bendtorch-core@0.1.0.0 · src/autograd/tensor_vjp_PROOF.bend:264

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/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/contracts/LAWS.bend as Contracts
import bendtorch-core@0.1.0.0/src/autograd/sum_to_size_fold_PROOF.bend as SumToSizeFoldProof
import bendtorch-core@0.1.0.0/src/tensor/layout.bend as Layout
import bendtorch-core@0.1.0.0/src/tensor/observation.bend as Values
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.bend as PermuteValues
import bendtorch-core@0.1.0.0/src/tensor/permute_values_records.bend as PermuteRecords

264# A permuted copy holds, at every in-bounds index, the source's word at
265# permute_source: its words are `permuted_words`, which are the words of the
266# copy loop's history.
267def permuted_at(+input: Shape.Shape, +fill: F32, +ops: List<&2, StorageLaws.Op>, +axes: List<&2, U32>, +plan: Layout.Plan<input, axes>,
268  +output_ops: List<&2, StorageLaws.Op>,
269  +moved: {Contracts.words(Layout.output(input, axes, plan), 0.0, output_ops, Shape.logical_numel(Shape.dimensions(Layout.output(input, axes, plan)))) ==
270    PermuteRecords.permuted_words(input, fill, ops, axes, Shape.logical_numel(Shape.dimensions(Layout.output(input, axes, plan)))) : List<&2, F32>},
271  +index: U32, +inside: {Nat.is_lt(U32.to_nat(index), Shape.logical_numel(Shape.dimensions(Layout.output(input, axes, plan)))) == True{} : Bool}) ->
272  {Replayed.value(Layout.output(input, axes, plan), 0.0, output_ops, index) ==
273    Replayed.value(input, fill, ops, U32.from_nat(Values.permute_source_index(input, axes, U32.to_nat(index)))) : F32}:
274  +output = Layout.output(input, axes, plan)
275  +total = Shape.logical_numel(Shape.dimensions(output))
276  +request : PermuteValues.Request = PermuteValues.Request{input, axes}
277  +history = IndexMap.writes(~PermuteValues.Request, ~PermuteValues.source_of, total, request, input, fill, ops)
278  Equal.trans(F32, Replayed.value(output, 0.0, output_ops, index), Replayed.value(output, 0.0, history, index),
279    Replayed.value(input, fill, ops, U32.from_nat(Values.permute_source_index(input, axes, U32.to_nat(index)))),
280    SumToSizeFoldProof.agree_at(output, 0.0, output_ops, history,
281      Equal.trans(List<&2, F32>, Contracts.words(output, 0.0, output_ops, total), PermuteRecords.permuted_words(input, fill, ops, axes, total),
282        Contracts.words(output, 0.0, history, total), moved,
283        Equal.sym(List<&2, F32>, Contracts.words(output, 0.0, history, total), PermuteRecords.permuted_words(input, fill, ops, axes, total),
284          PermuteValues.words_moved(input, fill, ops, axes, plan, total, Conversion.le_reflexive(total)))),
285      index, inside),
286    PointwiseProof.writes_value(~Unit, ~IndexMap.keep, ~PermuteValues.Request, ~PermuteValues.source_of, ~PermuteValues.source_of, total, Unit{},
287      request, input, fill, ops, input, fill, ops, output, index, inside, Conversion.le_reflexive(total)))