def PermutedWith

bendtorch-core@0.1.0.0 · src/tensor/permute_values_records.bend:55

import Base
import bendtorch-core@0.1.0.0/src/shape/core.bend as Shape
import bendtorch-core@0.1.0.0/src/tensor/core.bend as Core
import bendtorch-core@0.1.0.0/src/storage/LAWS.bend as StorageLaws
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/transforms.bend as Transforms
import bendtorch-core@0.1.0.0/src/tensor/contracts/replayed.bend as Replayed
import bendtorch-core@0.1.0.0/src/tensor/contracts/LAWS.bend as Contracts

55# What the transport form of permute_replayed and transpose_replayed asks of
56# P: P holds for the returned source, replayed with any history that keeps its
57# words, next to the output, of the plan's shape, replayed over +0.0 with any
58# history whose words are `permuted_words`.
59def PermutedWith(+input: Shape.Shape, +fill: F32, +ops: List<&2, StorageLaws.Op>, +axes: List<&2, U32>, +plan: Layout.Plan<input, axes>,
60  -P: Core.Dynamic & Result<&2, &1, Transforms.Error, Core.Dynamic> -> Type) -> Type:
61  @+kept: List<&2, StorageLaws.Op> ->
62  @kept_words: {Contracts.words(input, fill, kept, Shape.logical_numel(Shape.dimensions(input))) ==
63    Contracts.words(input, fill, ops, Shape.logical_numel(Shape.dimensions(input))) : List<&2, F32>} ->
64  @-kept_fits: {Lifecycle.logical_length(Replayed.storage(input, fill, kept)) == Shape.logical_numel(Shape.dimensions(input)) : Nat} ->
65  @+output_ops: List<&2, StorageLaws.Op> ->
66  @moved: {Contracts.words(Layout.output(input, axes, plan), 0.0, output_ops, Shape.logical_numel(Shape.dimensions(Layout.output(input, axes, plan)))) ==
67    permuted_words(input, fill, ops, axes, Shape.logical_numel(Shape.dimensions(Layout.output(input, axes, plan)))) : List<&2, F32>} ->
68  @-output_fits: {Lifecycle.logical_length(Replayed.storage(Layout.output(input, axes, plan), 0.0, output_ops)) ==
69    Shape.logical_numel(Shape.dimensions(Layout.output(input, axes, plan))) : Nat} ->
70    P((Core.Packed{input, Core.TTensor{Replayed.storage(input, fill, kept), kept_fits}},
71      Done{Core.Packed{Layout.output(input, axes, plan),
72        Core.TTensor{Replayed.storage(Layout.output(input, axes, plan), 0.0, output_ops), output_fits}}}))