1 import Base 2 import ../shape/core.bend as Shape 3 import ./core.bend as Core 4 import ./axis_spec.bend as Axis 5 import ./transpose_spec.bend as Transpose 6 import ./observation.bend as Values 7 import ./transform_observation.bend as Observation 8 import ../storage/LAWS.bend as StorageLaws 9 import ./lifecycle.bend as Lifecycle 10 import ./layout.bend as Layout 11 import ./transforms.bend as Transforms 12 import ./contracts/replayed.bend as Replayed 13 import ./contracts/LAWS.bend as Contracts 14 15 type PermuteContract<-input: Shape.Shape, -values: List<&2, F32>, -axes: List<&2, U32>> is Type: 16 PermuteContract{ 17 outcome: {Observation.failure_code(input, Observation.Permute{axes}, Core.from_values(input, values)) == Some{0} : Maybe<&2, U32>}, 18 dimensions: {Observation.applied_dimensions(input, Observation.Permute{axes}, Core.from_values(input, values)) == 19 Some{Axis.order(U32, Shape.dimensions(input), Shape.naturals(axes))} : Maybe<&2, List<&2, U32>>}, 20 count: {Observation.applied_length(input, Observation.Permute{axes}, Core.from_values(input, values)) == 21 Some{Shape.logical_numel(Shape.dimensions(input))} : Maybe<&2, Nat>}, 22 word_count: {Observation.applied_count(input, Observation.Permute{axes}, Core.from_values(input, values)) == Some{Shape.count(input)} : Maybe<&2, U32>}, 23 value: @+index: U32 -> @inside: {Nat.is_lt(U32.to_nat(index), Shape.logical_numel(Shape.dimensions(input))) == True{} : Bool} -> 24 {Observation.applied_value(input, Observation.Permute{axes}, index, Core.from_values(input, values)) == 25 Values.listed(values, Values.permute_source(input, axes, index)) : Maybe<&2, F32>}, 26 retained_value: @+other: U32 -> @inside: {Nat.is_lt(U32.to_nat(other), Shape.logical_numel(Shape.dimensions(input))) == True{} : Bool} -> 27 {Observation.kept_value(input, Observation.Permute{axes}, other, Core.from_values(input, values)) == 28 List.get(&2, F32, values, U32.to_nat(other)) : Maybe<&2, F32>}} 29 30 type TransposeContract<-input: Shape.Shape, -values: List<&2, F32>, -a: U32, -b: U32> is Type: 31 TransposeContract{ 32 outcome: {Observation.failure_code(input, Observation.Transpose{a, b}, Core.from_values(input, values)) == Some{0} : Maybe<&2, U32>}, 33 dimensions: {Observation.applied_dimensions(input, Observation.Transpose{a, b}, Core.from_values(input, values)) == 34 Some{Axis.order(U32, Shape.dimensions(input), Shape.naturals(Transpose.requested_axes(Shape.rank_dimensions(Shape.dimensions(input)), U32.to_nat(a), U32.to_nat(b))))} : Maybe<&2, List<&2, U32>>}, 35 count: {Observation.applied_length(input, Observation.Transpose{a, b}, Core.from_values(input, values)) == 36 Some{Shape.logical_numel(Shape.dimensions(input))} : Maybe<&2, Nat>}, 37 word_count: {Observation.applied_count(input, Observation.Transpose{a, b}, Core.from_values(input, values)) == Some{Shape.count(input)} : Maybe<&2, U32>}, 38 value: @+index: U32 -> @inside: {Nat.is_lt(U32.to_nat(index), Shape.logical_numel(Shape.dimensions(input))) == True{} : Bool} -> 39 {Observation.applied_value(input, Observation.Transpose{a, b}, index, Core.from_values(input, values)) == 40 Values.listed(values, Values.permute_source(input, Transpose.requested_axes(Shape.rank_dimensions(Shape.dimensions(input)), U32.to_nat(a), U32.to_nat(b)), index)) : Maybe<&2, F32>}, 41 retained_value: @+other: U32 -> @inside: {Nat.is_lt(U32.to_nat(other), Shape.logical_numel(Shape.dimensions(input))) == True{} : Bool} -> 42 {Observation.kept_value(input, Observation.Transpose{a, b}, other, Core.from_values(input, values)) == 43 List.get(&2, F32, values, U32.to_nat(other)) : Maybe<&2, F32>}} 44 45 # The words of a permuted copy in index order: word k is the source's word at 46 # `permute_source(input, axes, k)`. 47 def permuted_words(+input: Shape.Shape, +fill: F32, +ops: List<&2, StorageLaws.Op>, +axes: List<&2, U32>, +count: Nat) -> List<&2, F32>: 48 match count: 49 case 0n: 50 Nil{} 51 case 1n+previous: 52 List.append(&2, F32, permuted_words(input, fill, ops, axes, previous), 53 [Replayed.value(input, fill, ops, U32.from_nat(Values.permute_source_index(input, axes, previous)))]) 54 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`. 59 def PermutedWith(+input: Shape.Shape, +fill: F32, +ops: List<&2, StorageLaws.Op>, +axes: List<&2, U32>, +plan: Layout.Plan, 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}}}))