src/tensor/matmul_LAWS.bend
bendtorch-core@0.1.0.0 · 264 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/core.bend as Core
4import bendtorch-core@0.1.0.0/src/tensor/matmul.bend as Matrix
5import bendtorch-core@0.1.0.0/src/tensor/math_totality.bend as Totality
6import bendtorch-core@0.1.0.0/src/tensor/dot_fold.bend as DotFold
7
8def coordinate_observe_result(+shape: Shape.Shape, tensor: Core.Tensor<shape>,
9 result: Result<&2, &2, Shape.Error, Shape.FlatIndex<shape>>) -> Maybe<&2, F32>:
10 match result:
11 case Fail{error}:
12 None{}
13 case Done{index}:
14 Core.checked_observe(shape, tensor, Shape.index_value(shape, index))
15
16def coordinate_observe(+shape: Shape.Shape, tensor: Core.Tensor<shape>, coordinates: List<&2, U32>) -> Maybe<&2, F32>:
17 coordinate_observe_result(shape, tensor, Shape.flatten_checked(shape, coordinates))
18
19def LogicalDot(+fuel: Nat, +left: Shape.Shape, +right: Shape.Shape, -a: Core.Tensor<left>, -b: Core.Tensor<right>,
20 +row: U32, +column: U32, +cursor: Nat, values: List<&2, DotFold.Factors>) -> Data:
21 match fuel values:
22 case 0n Nil{}:
23 Unit
24 case 1n+tail DotFold.Factors{lv, rv} <> rest:
25 DotStep<tail, left, right, a, b, row, column, cursor, lv, rv, rest>
26 case 0n DotFold.Factors{lv, rv} <> rest:
27 Empty
28 case 1n+tail Nil{}:
29 Empty
30
31type DotStep<-fuel: Nat, -left: Shape.Shape, -right: Shape.Shape, -a: Core.Tensor<left>, -b: Core.Tensor<right>,
32 -row: U32, -column: U32, -cursor: Nat, -lv: F32, -rv: F32, -values: List<&2, DotFold.Factors>> is Data:
33 DotStep{-exact: {U32.to_nat(U32.from_nat(cursor)) == cursor : Nat},
34 -left_observed: {coordinate_observe(left, a, [row, U32.from_nat(cursor)]) == Some{lv} : Maybe<&2, F32>},
35 -right_observed: {coordinate_observe(right, b, [U32.from_nat(cursor), column]) == Some{rv} : Maybe<&2, F32>},
36 -remaining: LogicalDot(fuel, left, right, a, b, row, column, 1n+cursor, values)}
37
38type DotMeaning<-fuel: Nat, -left: Shape.Shape, -right: Shape.Shape, -a: Core.Tensor<left>, -b: Core.Tensor<right>,
39 -row: U32, -column: U32, -cursor: Nat, -accumulator: F32, -ra: Core.Tensor<left>, -rb: Core.Tensor<right>, -value: F32> is Data:
40 DotMeaning{-values: List<&2, DotFold.Factors>, -logical: LogicalDot(fuel, left, right, a, b, row, column, cursor, values),
41 -fold: {value == DotFold.dot_fold(values, accumulator) : F32},
42 -left_kept: @+index: U32 -> {Core.checked_observe(left, ra, index) == Core.checked_observe(left, a, index) : Maybe<&2, F32>},
43 -right_kept: @+index: U32 -> {Core.checked_observe(right, rb, index) == Core.checked_observe(right, b, index) : Maybe<&2, F32>}}
44
45def DotResultMeaning(-fuel: Nat, -left: Shape.Shape, -right: Shape.Shape, -a: Core.Tensor<left>, -b: Core.Tensor<right>,
46 -row: U32, -column: U32, -cursor: Nat, -accumulator: F32,
47 result: Result<&2, &1, Matrix.Error, Core.Tensor<left> & Core.Tensor<right> & F32>) -> Data:
48 match result:
49 case Fail{error}:
50 Unit
51 case Done{pair}:
52 (ra, rb, value) = pair
53 DotMeaning<fuel, left, right, a, b, row, column, cursor, accumulator, ra, rb, value>
54
55type SourceDot<-fuel: Nat, -left: Shape.Shape, -right: Shape.Shape, -a: Core.Tensor<left>, -b: Core.Tensor<right>,
56 -row: U32, -column: U32, -cursor: Nat, -accumulator: F32,
57 -result: Result<&2, &1, Matrix.Error, Core.Tensor<left> & Core.Tensor<right> & F32>> is Data:
58 SourceDot{-meaning: @unit: Unit -> DotResultMeaning(fuel, left, right, a, b, row, column, cursor, accumulator, result)}
59
60type ProductCell<-inner: U32, -left: Shape.Shape, -right: Shape.Shape, -a: Core.Tensor<left>, -b: Core.Tensor<right>,
61 -output: Shape.Shape, -out: Core.Tensor<output>, -position: Nat> is Data:
62 ProductCell{-row: U32, -column: U32, -coordinate: Shape.CheckedCoordinate<output>,
63 -decoded: {Shape.unflatten_checked(output, U32.from_nat(position)) == Done{coordinate} : Result<&2, &2, Shape.Error, Shape.CheckedCoordinate<output>>},
64 -coordinates: {Shape.coordinate_values(output, coordinate) == [row, column] : List<&2, U32>},
65 -word_exact: {U32.to_nat(U32.from_nat(position)) == position : Nat},
66 -factors: List<&2, DotFold.Factors>, -logical: LogicalDot(U32.to_nat(inner), left, right, a, b, row, column, 0n, factors),
67 -value: F32, -fold: {value == DotFold.dot_fold(factors, 0.0) : F32},
68 -checked: Core.ObservationChecked(Core.Packed{output, out}, U32.from_nat(position)),
69 -observed: {Core.checked_observe(output, out, U32.from_nat(position)) == Some{value} : Maybe<&2, F32>}}
70
71type ProductPrefix<-inner: U32, -left: Shape.Shape, -right: Shape.Shape, -a: Core.Tensor<left>, -b: Core.Tensor<right>,
72 -output: Shape.Shape, -ra: Core.Tensor<left>, -rb: Core.Tensor<right>, -out: Core.Tensor<output>, -cursor: Nat> is Data:
73 ProductPrefix{-left_kept: @+other: U32 -> {Core.checked_observe(left, ra, other) == Core.checked_observe(left, a, other) : Maybe<&2, F32>},
74 -right_kept: @+other: U32 -> {Core.checked_observe(right, rb, other) == Core.checked_observe(right, b, other) : Maybe<&2, F32>},
75 -cells: @+position: Nat -> @earlier: {Nat.is_lt(position, cursor) == True{} : Bool} -> ProductCell<inner, left, right, a, b, output, out, position>}
76
77def ProductMeaning(-inner: U32, -left: Shape.Shape, -right: Shape.Shape, -a: Core.Tensor<left>, -b: Core.Tensor<right>,
78 -output: Shape.Shape, -end: Nat, result: Result<&2, &1, Matrix.Error, Core.Tensor<left> & Core.Tensor<right> & Core.Tensor<output>>) -> Data:
79 match result:
80 case Fail{error}:
81 Unit
82 case Done{pair}:
83 (ra, rb, out) = pair
84 ProductPrefix<inner, left, right, a, b, output, ra, rb, out, end>
85
86type SourceProduct<-inner: U32, -left: Shape.Shape, -right: Shape.Shape, -a: Core.Tensor<left>, -b: Core.Tensor<right>,
87 -output: Shape.Shape, -end: Nat, -result: Result<&2, &1, Matrix.Error, Core.Tensor<left> & Core.Tensor<right> & Core.Tensor<output>>> is Data:
88 SourceProduct{-meaning: @unit: Unit -> ProductMeaning(inner, left, right, a, b, output, end, result)}
89
90law successful_product_interpretation:
91 for +fuel: Nat
92 for +left: Shape.Shape
93 for +right: Shape.Shape
94 for +output: Shape.Shape
95 for +inner: U32
96 for +columns: U32
97 for -a: Core.Tensor<left>
98 for -b: Core.Tensor<right>
99 for ra: Core.Tensor<left>
100 for rb: Core.Tensor<right>
101 for out: Core.Tensor<output>
102 for +cursor: Nat
103 for prefix: ProductPrefix<inner, left, right, a, b, output, ra, rb, out, cursor>
104 SourceProduct<inner, left, right, a, b, output, Nat.add(fuel, cursor),
105 Matrix.natural_product_loop(fuel, left, right, output, inner, columns, ra, rb, out, cursor)>
106
107def plan_inner(-left: Shape.Shape, -right: Shape.Shape, plan: Matrix.Plan<left, right>) -> U32:
108 Matrix.Plan{output, rows, inner, right_inner, columns, le, re, oe, compatible} = plan
109 inner
110
111law successful_multiply_interpretation:
112 for +left: Shape.Shape
113 for +right: Shape.Shape
114 for +plan: Matrix.Plan<left, right>
115 for a: Core.Tensor<left>
116 for b: Core.Tensor<right>
117 SourceProduct<plan_inner(left, right, plan), left, right, a, b, Matrix.output_shape(left, right, plan),
118 Shape.logical_numel(Shape.dimensions(Matrix.output_shape(left, right, plan))), Matrix.multiply(left, right, plan, a, b)>
119
120type MatmulSuccess<-original_a: Core.Dynamic, -original_b: Core.Dynamic, -retained_a: Core.Dynamic, -retained_b: Core.Dynamic, -product: Core.Dynamic> is Data:
121 MatmulSuccess{-left: Shape.Shape, -right: Shape.Shape, -plan: Matrix.Plan<left, right>,
122 -a: Core.Tensor<left>, -b: Core.Tensor<right>, -ra: Core.Tensor<left>, -rb: Core.Tensor<right>, -out: Core.Tensor<Matrix.output_shape(left, right, plan)>,
123 -input_left: {original_a == Core.Packed{left, a} : Core.Dynamic}, -input_right: {original_b == Core.Packed{right, b} : Core.Dynamic},
124 -output_left: {retained_a == Core.Packed{left, ra} : Core.Dynamic}, -output_right: {retained_b == Core.Packed{right, rb} : Core.Dynamic},
125 -output: {product == Core.Packed{Matrix.output_shape(left, right, plan), out} : Core.Dynamic},
126 -source: ProductPrefix<plan_inner(left, right, plan), left, right, a, b, Matrix.output_shape(left, right, plan), ra, rb, out,
127 Shape.logical_numel(Shape.dimensions(Matrix.output_shape(left, right, plan)))>}
128
129def MatmulMeaning(-a: Core.Dynamic, -b: Core.Dynamic, result: Result<&2, &1, Matrix.Error, Core.Dynamic & Core.Dynamic & Core.Dynamic>) -> Data:
130 match result:
131 case Fail{error}:
132 Unit
133 case Done{pair}:
134 (ra, rb, product) = pair
135 MatmulSuccess<a, b, ra, rb, product>
136
137type SourceMatmul<-a: Core.Dynamic, -b: Core.Dynamic, -result: Result<&2, &1, Matrix.Error, Core.Dynamic & Core.Dynamic & Core.Dynamic>> is Data:
138 SourceMatmul{-meaning: @unit: Unit -> MatmulMeaning(a, b, result)}
139
140law successful_matmul_interpretation:
141 for left: Core.Dynamic
142 for right: Core.Dynamic
143 SourceMatmul<left, right, Matrix.matmul(left, right)>
144
145type ZeroPrefix<-left: Shape.Shape, -right: Shape.Shape, -output: Shape.Shape, -out: Core.Tensor<output>, -end: Nat> is Data:
146 ZeroPrefix{-observed: @+position: Nat -> @earlier: {Nat.is_lt(position, end) == True{} : Bool} ->
147 {Core.checked_observe(output, out, U32.from_nat(position)) == Some{0.0} : Maybe<&2, F32>}}
148
149def ZeroMeaning(-left: Shape.Shape, -right: Shape.Shape, -output: Shape.Shape, -end: Nat,
150 result: Result<&2, &1, Matrix.Error, Core.Tensor<left> & Core.Tensor<right> & Core.Tensor<output>>) -> Data:
151 match result:
152 case Fail{error}:
153 Unit
154 case Done{pair}:
155 (a, b, out) = pair
156 ZeroPrefix<left, right, output, out, end>
157
158type ZeroResult<-left: Shape.Shape, -right: Shape.Shape, -output: Shape.Shape, -end: Nat,
159 -result: Result<&2, &1, Matrix.Error, Core.Tensor<left> & Core.Tensor<right> & Core.Tensor<output>>> is Data:
160 ZeroResult{-meaning: @unit: Unit -> ZeroMeaning(left, right, output, end, result)}
161
162law zero_inner_multiply:
163 for +left: Shape.Shape
164 for +right: Shape.Shape
165 for +output: Shape.Shape
166 for +rows: U32
167 for +right_inner: U32
168 for +columns: U32
169 for +le: {Shape.dimensions(left) == [rows, 0] : List<&2, U32>}
170 for +re: {Shape.dimensions(right) == [right_inner, columns] : List<&2, U32>}
171 for +oe: {Shape.dimensions(output) == [rows, columns] : List<&2, U32>}
172 for +compatible: {0n == U32.to_nat(right_inner) : Nat}
173 for a: Core.Tensor<left>
174 for b: Core.Tensor<right>
175 ZeroResult<left, right, output, Shape.logical_numel(Shape.dimensions(output)),
176 Matrix.multiply(left, right, Matrix.Plan{output, rows, 0, right_inner, columns, le, re, oe, compatible}, a, b)>
177
178law successful_dot_interpretation:
179 for +fuel: Nat
180 for +left: Shape.Shape
181 for +right: Shape.Shape
182 for a: Core.Tensor<left>
183 for b: Core.Tensor<right>
184 for +row: U32
185 for +column: U32
186 for +cursor: Nat
187 for +accumulator: F32
188 SourceDot<fuel, left, right, a, b, row, column, cursor, accumulator,
189 Matrix.natural_dot(fuel, left, right, a, b, row, column, cursor, accumulator)>
190
191law zero_inner_dot:
192 for +left: Shape.Shape
193 for +right: Shape.Shape
194 for a: Core.Tensor<left>
195 for b: Core.Tensor<right>
196 for +row: U32
197 for +column: U32
198 {Matrix.natural_dot(0n, left, right, a, b, row, column, 0n, 0.0) == Done{(a, b, 0.0)} :
199 Result<&2, &1, Matrix.Error, Core.Tensor<left> & Core.Tensor<right> & F32>}
200
201law output_dimensions:
202 for -left: Shape.Shape
203 for -right: Shape.Shape
204 for +output: Shape.Shape
205 for rows: U32
206 for inner: U32
207 for right_inner: U32
208 for columns: U32
209 for le: {Shape.dimensions(left) == [rows, inner] : List<&2, U32>}
210 for re: {Shape.dimensions(right) == [right_inner, columns] : List<&2, U32>}
211 for oe: {Shape.dimensions(output) == [rows, columns] : List<&2, U32>}
212 for compatible: {U32.to_nat(inner) == U32.to_nat(right_inner) : Nat}
213 {Shape.dimensions(Matrix.output_shape(left, right, Matrix.Plan{output, rows, inner, right_inner, columns, le, re, oe, compatible})) ==
214 [rows, columns] : List<&2, U32>}
215
216type TotalMatmul<-left: Core.Dynamic, -right: Core.Dynamic,
217 -result: Result<&2, &1, Matrix.Error, Core.Dynamic & Core.Dynamic & Core.Dynamic>> is Data:
218 TotalMatmul{-done: Totality.Success<Matrix.Error, Core.Dynamic & Core.Dynamic & Core.Dynamic, result>,
219 -interpretation: SourceMatmul<left, right, result>}
220
221law matmul_valid_success:
222 for +left: Shape.Shape
223 for +right: Shape.Shape
224 for a: Core.Tensor<left>
225 for b: Core.Tensor<right>
226 for valid: Totality.MatrixInputs<left, right>
227 TotalMatmul<Core.Packed{left, a}, Core.Packed{right, b}, Matrix.matmul(Core.Packed{left, a}, Core.Packed{right, b})>
228
229
230
231law matmul_left_rank_error:
232 for +left: Shape.Shape
233 for +right: Shape.Shape
234 for a: Core.Tensor<left>
235 for b: Core.Tensor<right>
236 for not_matrix: {Shape.rank_dimensions(Shape.dimensions(left)) != 2n : Nat}
237 {Matrix.matmul(Core.Packed{left, a}, Core.Packed{right, b}) == Fail{Matrix.RankMismatch{}} : Result<&2, &1, Matrix.Error, Core.Dynamic & Core.Dynamic & Core.Dynamic>}
238
239
240
241law matmul_right_rank_error:
242 for +left: Shape.Shape
243 for +right: Shape.Shape
244 for a: Core.Tensor<left>
245 for b: Core.Tensor<right>
246 for not_matrix: {Shape.rank_dimensions(Shape.dimensions(right)) != 2n : Nat}
247 {Matrix.matmul(Core.Packed{left, a}, Core.Packed{right, b}) == Fail{Matrix.RankMismatch{}} : Result<&2, &1, Matrix.Error, Core.Dynamic & Core.Dynamic & Core.Dynamic>}
248
249
250
251
252law matmul_inner_error:
253 for +left: Shape.Shape
254 for +right: Shape.Shape
255 for a: Core.Tensor<left>
256 for b: Core.Tensor<right>
257 for +rows: U32
258 for +inner: U32
259 for +right_inner: U32
260 for +columns: U32
261 for left_matrix: {Shape.dimensions(left) == [rows, inner] : List<&2, U32>}
262 for right_matrix: {Shape.dimensions(right) == [right_inner, columns] : List<&2, U32>}
263 for mismatch: {U32.to_nat(inner) != U32.to_nat(right_inner) : Nat}
264 {Matrix.matmul(Core.Packed{left, a}, Core.Packed{right, b}) == Fail{Matrix.InnerDimensionMismatch{}} : Result<&2, &1, Matrix.Error, Core.Dynamic & Core.Dynamic & Core.Dynamic>}