1 import Base 2 import ../shape/core.bend as Shape 3 import ./layout.bend as Layout 4 import ./transpose_spec.bend as Transpose 5 import ./permute_values_records.bend as Records 6 import ../storage/LAWS.bend as StorageLaws 7 import ./core.bend as Core 8 import ./lifecycle.bend as Lifecycle 9 import ./transforms.bend as Transforms 10 import ./contracts/replayed.bend as Replayed 11 12 # `permute_copy` of a `Core.from_values(input, values)` tensor by a valid 13 # permutation succeeds with dimensions `Axis.order(dims, axes)`; output index i 14 # holds `values[permute_source(input, axes, i)]`, and the returned source 15 # still holds `values[j]` at every j. 16 law permute_contract: 17 for +input: Shape.Shape 18 for +values: List<&2, F32> 19 for counted: {List.length(&2, F32, values) == Shape.logical_numel(Shape.dimensions(input)) : Nat} 20 for +axes: List<&2, U32> 21 for valid: {True{} == Layout.bijective(input, axes) : Bool} 22 Records.PermuteContract 23 24 # `transpose_copy(a, b)` of a `Core.from_values` tensor with both axes in 25 # range is the permute contract with the axes that swap a and b. 26 law transpose_contract: 27 for +input: Shape.Shape 28 for +values: List<&2, F32> 29 for counted: {List.length(&2, F32, values) == Shape.logical_numel(Shape.dimensions(input)) : Nat} 30 for +a: U32 31 for +b: U32 32 for inside: {True{} == Transpose.valid(input, a, b) : Bool} 33 Records.TransposeContract 34 35 # `permute_dynamic` of a replayed tensor, once `Layout.permutation` gives the 36 # plan, in the transport form of `broadcast_to_replayed`: for every P, if P 37 # holds for every result that hands the source back with the same words and 38 # gives an output of the plan's shape, replayed over +0.0, whose cell i holds 39 # the source's word at `permute_source(input, axes, i)`, then P holds for the 40 # call. The output is a replayed tensor, so a later family law applies to it. 41 law permute_replayed: 42 for +input: Shape.Shape 43 for +fill: F32 44 for +ops: List<&2, StorageLaws.Op> 45 for -fits: {Lifecycle.logical_length(Replayed.storage(input, fill, ops)) == Shape.logical_numel(Shape.dimensions(input)) : Nat} 46 for +axes: List<&2, U32> 47 for +plan: Layout.Plan 48 for laid: {Layout.permutation(input, axes) == Done{plan} : Result<&2, &2, Layout.Error, Layout.Plan>} 49 for -P: Core.Dynamic & Result<&2, &1, Transforms.Error, Core.Dynamic> -> Type 50 for h: Records.PermutedWith(input, fill, ops, axes, plan, P) 51 P(Transforms.permute_dynamic(Core.Packed{input, Core.TTensor{Replayed.storage(input, fill, ops), fits}}, axes)) 52 53 # `transpose_dynamic(a, b)` of a replayed tensor with both axes in range is the 54 # permute of `permute_replayed` by the swapped axes, in the same transport form. 55 law transpose_replayed: 56 for +input: Shape.Shape 57 for +fill: F32 58 for +ops: List<&2, StorageLaws.Op> 59 for -fits: {Lifecycle.logical_length(Replayed.storage(input, fill, ops)) == Shape.logical_numel(Shape.dimensions(input)) : Nat} 60 for +a: U32 61 for +b: U32 62 for inside: {True{} == Transpose.valid(input, a, b) : Bool} 63 for +plan: Layout.Plan 64 for laid: {Layout.permutation(input, Transpose.swapped_axes(input, a, b)) == Done{plan} : 65 Result<&2, &2, Layout.Error, Layout.Plan>} 66 for -P: Core.Dynamic & Result<&2, &1, Transforms.Error, Core.Dynamic> -> Type 67 for h: Records.PermutedWith(input, fill, ops, Transpose.swapped_axes(input, a, b), plan, P) 68 P(Transforms.transpose_dynamic(Core.Packed{input, Core.TTensor{Replayed.storage(input, fill, ops), fits}}, a, b))