src/tensor/transform_totality.bend
bendtorch-core@0.1.0.0 · 164 lines · 30 definitions
1import Base
2import bendtorch-core@0.1.0.0/src/shape/core.bend as Shape
3import bendtorch-core@0.1.0.0/src/tensor/layout.bend as Layout
4import bendtorch-core@0.1.0.0/src/tensor/core.bend as Core
5import bendtorch-core@0.1.0.0/src/tensor/transforms.bend as Transforms
6import bendtorch-core@0.1.0.0/src/tensor/axis_spec.bend as Axis
7import bendtorch-core@0.1.0.0/src/tensor/transpose_spec.bend as Transpose
8
9type PlanSuccess<-input: Shape.Shape, -axes: List<&2, U32>> is Data:
10 PlanSuccess{-successful: {Result.is_done(&2, &2, Layout.Error, Layout.Plan<input, axes>,
11 Layout.permutation(input, axes)) == True{} : Bool}}
12
13def read_succeeded(-shape: Shape.Shape, pair: Core.Tensor<shape> & Result<&2, &2, Core.Error, F32>) -> Bool:
14 (owner, result) = pair
15 Result.is_done(&2, &2, Core.Error, F32, result)
16
17type ReadSuccess<-shape: Shape.Shape, -original: Core.Tensor<shape>, -index: U32> is Data:
18 ReadSuccess{-successful: {read_succeeded(shape, Core.read_flat(shape, original, index)) == True{} : Bool}}
19
20def write_succeeded(-shape: Shape.Shape, pair: Core.Tensor<shape> & Result<&2, &2, Core.Error, Unit>) -> Bool:
21 (owner, result) = pair
22 Result.is_done(&2, &2, Core.Error, Unit, result)
23
24type WriteSuccess<-shape: Shape.Shape, -original: Core.Tensor<shape>, -index: U32, -value: F32> is Data:
25 WriteSuccess{-successful: {write_succeeded(shape, Core.write_flat(shape, original, index, value)) == True{} : Bool}}
26
27type WordsSuccess<-dimensions: List<&2, U32>, -values: List<&2, U32>> is Data:
28 WordsSuccess{-successful: {Result.is_done(&2, &2, Shape.Error, Layout.WordsValid(dimensions, values),
29 Layout.words_check(dimensions, values)) == True{} : Bool}}
30
31type PermuteFailure<-input: Shape.Shape, -axes: List<&2, U32>, -original: Core.Tensor<input>, -error: Layout.Error> is Data:
32 PermuteFailure{-exact: {Transforms.permute_copy(input, original, axes) == (original, Fail{Transforms.LayoutError{error}}) :
33 Core.Tensor<input> & Result<&2, &1, Transforms.Error, Core.Dynamic>}}
34
35type TransposeFailure<-input: Shape.Shape, -original: Core.Tensor<input>, -a: U32, -b: U32> is Data:
36 TransposeFailure{-exact: {Transforms.transpose_copy(input, original, a, b) == (original, Fail{Transforms.LayoutError{Layout.AxisOutOfBounds{}}}) :
37 Core.Tensor<input> & Result<&2, &1, Transforms.Error, Core.Dynamic>}}
38
39type ReshapeFailure<-input: Shape.Shape, -target: Shape.Shape, -original: Core.Tensor<input>> is Data:
40 ReshapeFailure{-exact: {Transforms.reshape(input, original, Shape.dimensions(target)) == Fail{Transforms.LayoutError{Layout.ElementCountMismatch{}}} :
41 Result<&2, &1, Transforms.Error, Core.Dynamic>}}
42
43type MappingSuccess<-input: Shape.Shape, -axes: List<&2, U32>, -plan: Layout.Plan<input, axes>, -index: U32> is Data:
44 MappingSuccess{-successful: {Result.is_done(&2, &2, Layout.Error, Layout.Mapping<input, axes, plan, index>,
45 Layout.map_index(input, axes, plan, index)) == True{} : Bool}}
46
47def copy_succeeded(-input: Shape.Shape, -axes: List<&2, U32>, -plan: Layout.Plan<input, axes>,
48 pair: Core.Tensor<input> & Result<&2, &1, Transforms.Error, Core.Tensor<Layout.output(input, axes, plan)>>) -> Bool:
49 (source, result) = pair
50 Result.is_done(&2, &1, Transforms.Error, Core.Tensor<Layout.output(input, axes, plan)>, result)
51
52type CopyStepSuccess<-input: Shape.Shape, -axes: List<&2, U32>, -plan: Layout.Plan<input, axes>,
53 -source: Core.Tensor<input>, -target: Core.Tensor<Layout.output(input, axes, plan)>, -index: U32> is Data:
54 CopyStepSuccess{-successful: {copy_succeeded(input, axes, plan, Transforms.copy_step(input, axes, plan, source, target, index)) == True{} : Bool}}
55
56type CopySuccess<-input: Shape.Shape, -axes: List<&2, U32>, -plan: Layout.Plan<input, axes>,
57 -pair: Core.Tensor<input> & Result<&2, &1, Transforms.Error, Core.Tensor<Layout.output(input, axes, plan)>>> is Data:
58 CopySuccess{-successful: {copy_succeeded(input, axes, plan, pair) == True{} : Bool}}
59
60def natural_succeeded(-input: Shape.Shape, -axes: List<&2, U32>, -plan: Layout.Plan<input, axes>, -cursor: Nat,
61 result: Transforms.NaturalCopy<input, axes, plan, cursor>) -> Bool:
62 match result:
63 case Transforms.NaturalError{source}:
64 False{}
65 case Transforms.NaturalCopied{pair, index, exact, candidate}:
66 copy_succeeded(input, axes, plan, pair)
67
68type NaturalSuccess<-input: Shape.Shape, -axes: List<&2, U32>, -plan: Layout.Plan<input, axes>, -cursor: Nat,
69 -result: Transforms.NaturalCopy<input, axes, plan, cursor>> is Data:
70 NaturalSuccess{-successful: {natural_succeeded(input, axes, plan, cursor, result) == True{} : Bool}}
71
72def public_succeeded(-input: Shape.Shape, pair: Core.Tensor<input> & Result<&2, &1, Transforms.Error, Core.Dynamic>) -> Bool:
73 (source, result) = pair
74 Result.is_done(&2, &1, Transforms.Error, Core.Dynamic, result)
75
76type PublicSuccess<-input: Shape.Shape, -pair: Core.Tensor<input> & Result<&2, &1, Transforms.Error, Core.Dynamic>> is Data:
77 PublicSuccess{-successful: {public_succeeded(input, pair) == True{} : Bool}}
78
79def inverse_axes(+input: Shape.Shape, axes: List<&2, U32>) -> List<&2, U32>:
80 Shape.words(Axis.inverse_axes(Shape.rank_dimensions(Shape.dimensions(input)), Shape.naturals(axes)))
81
82type TypedPermutation<-input: Shape.Shape, -axes: List<&2, U32>, -original: Core.Tensor<input>,
83 -pair: Core.Tensor<input> & Result<&2, &1, Transforms.Error, Core.Dynamic>> is Data:
84 TypedPermutation{-plan: Layout.Plan<input, axes>, -retained: Core.Tensor<input>, -target: Core.Tensor<Layout.output(input, axes, plan)>,
85 -exact: {(retained, Done{Core.Packed{Layout.output(input, axes, plan), target}}) == pair : Core.Tensor<input> & Result<&2, &1, Transforms.Error, Core.Dynamic>},
86 -prefix: Transforms.Prefix<input, axes, plan, original, retained, target, Shape.logical_numel(Shape.dimensions(Layout.output(input, axes, plan)))>}
87
88type MappingDone<-input: Shape.Shape, -axes: List<&2, U32>, -plan: Layout.Plan<input, axes>, -index: U32,
89 -result: Result<&2, &2, Layout.Error, Layout.Mapping<input, axes, plan, index>>> is Data:
90 MappingDone{-mapping: Layout.Mapping<input, axes, plan, index>,
91 -exact: {Done{mapping} == result : Result<&2, &2, Layout.Error, Layout.Mapping<input, axes, plan, index>>}}
92
93type CopyObservation<-input: Shape.Shape, -axes: List<&2, U32>, -plan: Layout.Plan<input, axes>, -original: Core.Tensor<input>,
94 -target: Core.Tensor<Layout.output(input, axes, plan)>, -index: U32> is Data:
95 CopyObservation{-equal: {Core.observe(Layout.output(input, axes, plan), target, index) == Transforms.expected_output(input, axes, plan, original, index) : Maybe<&2, F32>}}
96
97type RestoreObservation<-input: Shape.Shape, -output: Shape.Shape, -original: Core.Tensor<input>, -restored: Core.Tensor<output>, -index: U32> is Data:
98 RestoreObservation{-equal: {Core.observe(output, restored, index) == Core.observe(input, original, index) : Maybe<&2, F32>}}
99
100type PermutationRestoration<-input: Shape.Shape, -axes: List<&2, U32>, -original: Core.Tensor<input>> is Data:
101 PermutationRestoration{-first: Layout.Plan<input, axes>, -retained: Core.Tensor<input>, -middle: Core.Tensor<Layout.output(input, axes, first)>,
102 -first_exact: {(retained, Done{Core.Packed{Layout.output(input, axes, first), middle}}) == Transforms.permute_copy(input, original, axes) :
103 Core.Tensor<input> & Result<&2, &1, Transforms.Error, Core.Dynamic>},
104 -second: Layout.Plan<Layout.output(input, axes, first), inverse_axes(input, axes)>,
105 -middle_retained: Core.Tensor<Layout.output(input, axes, first)>,
106 -restored: Core.Tensor<Layout.output(Layout.output(input, axes, first), inverse_axes(input, axes), second)>,
107 -second_exact: {(Core.Packed{Layout.output(input, axes, first), middle_retained}, Done{Core.Packed{Layout.output(Layout.output(input, axes, first), inverse_axes(input, axes), second), restored}}) ==
108 Transforms.permute_dynamic(Core.Packed{Layout.output(input, axes, first), middle}, inverse_axes(input, axes)) : Core.Dynamic & Result<&2, &1, Transforms.Error, Core.Dynamic>},
109 -dimensions: {Shape.dimensions(Layout.output(Layout.output(input, axes, first), inverse_axes(input, axes), second)) == Shape.dimensions(input) : List<&2, U32>},
110 -linear: @+index: U32 -> @-inside: {Nat.is_lt(U32.to_nat(index), Shape.logical_numel(Shape.dimensions(input))) == True{} : Bool} ->
111 {Core.observe(Layout.output(Layout.output(input, axes, first), inverse_axes(input, axes), second), restored, index) == Core.observe(input, original, index) : Maybe<&2, F32>}}
112
113type PermutationInverse<-input: Shape.Shape, -axes: List<&2, U32>, -original: Core.Tensor<input>> is Data:
114 PermutationInverse{-restoration: PermutationRestoration<input, axes, original>}
115
116def transpose_reference(+input: Shape.Shape, +output: Shape.Shape, +a: U32, +b: U32, index: U32) -> Maybe<&2, U32>:
117 Layout.reference_source_coordinates(Shape.dimensions(input),
118 Axis.scatter(U32, Shape.rank_dimensions(Shape.dimensions(input)),
119 Shape.naturals(Transpose.requested_axes(Shape.rank_dimensions(Shape.dimensions(input)), U32.to_nat(a), U32.to_nat(b))),
120 Shape.words(Shape.digits(Shape.extents(output), U32.to_nat(index))), 0n))
121
122def reference_observe(-input: Shape.Shape, source: Core.Tensor<input>, reference: Maybe<&2, U32>) -> Maybe<&2, F32>:
123 match reference:
124 case None{}:
125 None{}
126 case Some{index}:
127 Core.observe(input, source, index)
128
129type TransposeObservation<-input: Shape.Shape, -output: Shape.Shape, -source: Core.Tensor<input>, -target: Core.Tensor<output>, -a: U32, -b: U32, -index: U32> is Data:
130 TransposeObservation{-equal: {Core.observe(output, target, index) == reference_observe(input, source, transpose_reference(input, output, a, b, index)) : Maybe<&2, F32>}}
131
132type RequestedTranspose<-input: Shape.Shape, -original: Core.Tensor<input>, -a: U32, -b: U32> is Data:
133 RequestedTranspose{-plan: Layout.Plan<input, Layout.transpose_axes(Shape.dimensions(input), a, b, 0)>,
134 -retained: Core.Tensor<input>, -target: Core.Tensor<Layout.output(input, Layout.transpose_axes(Shape.dimensions(input), a, b, 0), plan)>,
135 -exact: {(retained, Done{Core.Packed{Layout.output(input, Layout.transpose_axes(Shape.dimensions(input), a, b, 0), plan), target}}) ==
136 Transforms.transpose_copy(input, original, a, b) : Core.Tensor<input> & Result<&2, &1, Transforms.Error, Core.Dynamic>},
137 -dimensions: {Shape.dimensions(Layout.output(input, Layout.transpose_axes(Shape.dimensions(input), a, b, 0), plan)) ==
138 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)))) : List<&2, U32>},
139 -prefix: Transforms.Prefix<input, Layout.transpose_axes(Shape.dimensions(input), a, b, 0), plan, original, retained, target,
140 Shape.logical_numel(Shape.dimensions(Layout.output(input, Layout.transpose_axes(Shape.dimensions(input), a, b, 0), plan)))>,
141 -mapping: @+index: U32 -> @mapping: Layout.Mapping<input, Layout.transpose_axes(Shape.dimensions(input), a, b, 0), plan, index> ->
142 {Some{Layout.source_index(input, Layout.transpose_axes(Shape.dimensions(input), a, b, 0), plan, index, mapping)} ==
143 transpose_reference(input, Layout.output(input, Layout.transpose_axes(Shape.dimensions(input), a, b, 0), plan), a, b, index) : Maybe<&2, U32>},
144 -linear: @+index: U32 -> @-inside: {Nat.is_lt(U32.to_nat(index), Shape.logical_numel(Shape.dimensions(Layout.output(input, Layout.transpose_axes(Shape.dimensions(input), a, b, 0), plan)))) == True{} : Bool} ->
145 {Core.observe(Layout.output(input, Layout.transpose_axes(Shape.dimensions(input), a, b, 0), plan), target, index) ==
146 reference_observe(input, original, transpose_reference(input, Layout.output(input, Layout.transpose_axes(Shape.dimensions(input), a, b, 0), plan), a, b, index)) : Maybe<&2, F32>}}
147
148type TransposeRestoration<-input: Shape.Shape, -original: Core.Tensor<input>, -a: U32, -b: U32> is Data:
149 TransposeRestoration{-first: Layout.Plan<input, Layout.transpose_axes(Shape.dimensions(input), a, b, 0)>,
150 -retained: Core.Tensor<input>, -middle: Core.Tensor<Layout.output(input, Layout.transpose_axes(Shape.dimensions(input), a, b, 0), first)>,
151 -first_exact: {(retained, Done{Core.Packed{Layout.output(input, Layout.transpose_axes(Shape.dimensions(input), a, b, 0), first), middle}}) ==
152 Transforms.transpose_copy(input, original, a, b) : Core.Tensor<input> & Result<&2, &1, Transforms.Error, Core.Dynamic>},
153 -second: Layout.Plan<Layout.output(input, Layout.transpose_axes(Shape.dimensions(input), a, b, 0), first), inverse_axes(input, Layout.transpose_axes(Shape.dimensions(input), a, b, 0))>,
154 -middle_retained: Core.Tensor<Layout.output(input, Layout.transpose_axes(Shape.dimensions(input), a, b, 0), first)>,
155 -restored: Core.Tensor<Layout.output(Layout.output(input, Layout.transpose_axes(Shape.dimensions(input), a, b, 0), first), inverse_axes(input, Layout.transpose_axes(Shape.dimensions(input), a, b, 0)), second)>,
156 -second_exact: {(Core.Packed{Layout.output(input, Layout.transpose_axes(Shape.dimensions(input), a, b, 0), first), middle_retained},
157 Done{Core.Packed{Layout.output(Layout.output(input, Layout.transpose_axes(Shape.dimensions(input), a, b, 0), first), inverse_axes(input, Layout.transpose_axes(Shape.dimensions(input), a, b, 0)), second), restored}}) ==
158 Transforms.transpose_dynamic(Core.Packed{Layout.output(input, Layout.transpose_axes(Shape.dimensions(input), a, b, 0), first), middle}, a, b) : Core.Dynamic & Result<&2, &1, Transforms.Error, Core.Dynamic>},
159 -dimensions: {Shape.dimensions(Layout.output(Layout.output(input, Layout.transpose_axes(Shape.dimensions(input), a, b, 0), first), inverse_axes(input, Layout.transpose_axes(Shape.dimensions(input), a, b, 0)), second)) == Shape.dimensions(input) : List<&2, U32>},
160 -linear: @+index: U32 -> @-inside: {Nat.is_lt(U32.to_nat(index), Shape.logical_numel(Shape.dimensions(input))) == True{} : Bool} ->
161 {Core.observe(Layout.output(Layout.output(input, Layout.transpose_axes(Shape.dimensions(input), a, b, 0), first), inverse_axes(input, Layout.transpose_axes(Shape.dimensions(input), a, b, 0)), second), restored, index) == Core.observe(input, original, index) : Maybe<&2, F32>}}
162
163type TransposeInverse<-input: Shape.Shape, -original: Core.Tensor<input>, -a: U32, -b: U32> is Data:
164 TransposeInverse{-restoration: TransposeRestoration<input, original, a, b>}