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))