src/tensor/signed_LAWS.bend

bendtorch-core@0.1.0.0 · 412 lines · 51 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/layout.bend as Layout
5import bendtorch-core@0.1.0.0/src/tensor/transforms.bend as Transforms
6import bendtorch-core@0.1.0.0/src/tensor/transpose_spec.bend as Transpose
7import bendtorch-core@0.1.0.0/src/tensor/elementwise_totality.bend as Totality
8import bendtorch-core@0.1.0.0/src/tensor/signed.bend as Signed
9import bendtorch-core@0.1.0.0/src/tensor/slice.bend as Slice
10
11# PyTorch's signed axes, indices and -1 reshape sizes (signed.bend). These
12# laws state the wrap against PyTorch's definition, the range it accepts, the
13# reshape size it infers and the cases it refuses, and that each signed
14# operation is the unsigned one on the wrapped arguments; the unsigned
15# operations' laws then give the result and the refusal codes.
16
17# Exact sizes, the reshape request of the unsigned reshape.
18def exacts(dimensions: List<&2, U32>) -> List<&2, Signed.Extent>:
19  match dimensions:
20    case Nil{}:
21      Nil{}
22    case extent <> rest:
23      Signed.Exact{extent} <> exacts(rest)
24
25# A position counted from the start is kept.
26law signed_first_kept:
27  for +limit: U32
28  for +k: U32
29  {Signed.wrap(limit, Signed.First{k}) == k : U32}
30
31# Last{k} is the position d = -1 - k. In range, it wraps to d + limit:
32# the wrapped position plus k + 1 is limit, in exact arithmetic.
33law signed_last_inside:
34  for +limit: U32
35  for +k: U32
36  for inside: {U32.is_lt(k, limit) == True{} : Bool}
37  {Nat.add(U32.to_nat(Signed.wrap(limit, Signed.Last{k})), 1n+U32.to_nat(k)) == U32.to_nat(limit) : Nat}
38
39# Below -limit, it becomes limit itself, which every check refuses.
40law signed_last_outside:
41  for +limit: U32
42  for +k: U32
43  for outside: {U32.is_lt(k, limit) == False{} : Bool}
44  {Signed.wrap(limit, Signed.Last{k}) == limit : U32}
45
46# A wrapped position is below limit exactly when the position is in
47# PyTorch's range -limit..limit-1.
48law signed_wrap_inside:
49  for +limit: U32
50  for +position: Signed.Position
51  {U32.is_lt(Signed.wrap(limit, position), limit) == Signed.inside(limit, position) : Bool}
52
53# So a list of wrapped axes passes the operations' range check
54# (Layout.bounded) exactly when every signed axis is in range.
55law signed_axes_bounded:
56  for +limit: U32
57  for +positions: List<&2, Signed.Position>
58  {Layout.bounded(Signed.wrap_all(limit, positions), limit) == Signed.all_inside(limit, positions) : Bool}
59
60# Axes are wrapped one by one, in order, among the same limit.
61law signed_wrap_all_empty:
62  for +limit: U32
63  {Signed.wrap_all(limit, []) == [] : List<&2, U32>}
64
65law signed_wrap_all_step:
66  for +limit: U32
67  for +head: Signed.Position
68  for +tail: List<&2, Signed.Position>
69  {Signed.wrap_all(limit, head <> tail) == Signed.wrap(limit, head) <> Signed.wrap_all(limit, tail) : List<&2, U32>}
70
71# Coordinates are wrapped one by one, each among its own axis's extent; a
72# coordinate past the rank among 0.
73law signed_wrap_each_step:
74  for +extent: U32
75  for +rest: List<&2, U32>
76  for +head: Signed.Position
77  for +tail: List<&2, Signed.Position>
78  {Signed.wrap_each(extent <> rest, head <> tail) == Signed.wrap(extent, head) <> Signed.wrap_each(rest, tail) : List<&2, U32>}
79
80law signed_wrap_each_end:
81  for +dimensions: List<&2, U32>
82  {Signed.wrap_each(dimensions, []) == [] : List<&2, U32>}
83
84law signed_wrap_each_past_rank:
85  for +positions: List<&2, Signed.Position>
86  {Signed.wrap_each([], positions) == Signed.wrap_all(0, positions) : List<&2, U32>}
87
88# Transpose wraps among the rank, and among 1 for a rank-0 input.
89law signed_axis_limit:
90  for +input: Shape.Shape
91  {Signed.axis_limit(input) == U32.from_nat(Transpose.axis_limit(Shape.rank_dimensions(Shape.dimensions(input)))) : U32}
92
93# Each signed operation is the unsigned one on the wrapped arguments.
94law signed_permute_wraps:
95  for +input: Shape.Shape
96  for source: Core.Tensor<input>
97  for +axes: List<&2, Signed.Position>
98  {Signed.permute_copy(input, source, axes) == Transforms.permute_copy(input, source, Signed.wrap_all(Shape.rank(input), axes)) :
99    Core.Tensor<input> & Result<&2, &1, Transforms.Error, Core.Dynamic>}
100
101law signed_permute_dynamic_wraps:
102  for input: Core.Dynamic
103  for +axes: List<&2, Signed.Position>
104  {Signed.permute_dynamic(input, axes) == Transforms.permute_dynamic(input, Signed.wrap_all(Shape.rank(Totality.dynamic_shape(input)), axes)) :
105    Core.Dynamic & Result<&2, &1, Transforms.Error, Core.Dynamic>}
106
107law signed_transpose_wraps:
108  for +input: Shape.Shape
109  for source: Core.Tensor<input>
110  for +a: Signed.Position
111  for +b: Signed.Position
112  {Signed.transpose_copy(input, source, a, b) ==
113    Transforms.transpose_copy(input, source, Signed.wrap(Signed.axis_limit(input), a), Signed.wrap(Signed.axis_limit(input), b)) :
114    Core.Tensor<input> & Result<&2, &1, Transforms.Error, Core.Dynamic>}
115
116law signed_transpose_dynamic_wraps:
117  for input: Core.Dynamic
118  for +a: Signed.Position
119  for +b: Signed.Position
120  {Signed.transpose_dynamic(input, a, b) ==
121    Transforms.transpose_dynamic(input, Signed.wrap(Signed.axis_limit(Totality.dynamic_shape(input)), a),
122      Signed.wrap(Signed.axis_limit(Totality.dynamic_shape(input)), b)) :
123    Core.Dynamic & Result<&2, &1, Transforms.Error, Core.Dynamic>}
124
125law signed_read_coordinate_wraps:
126  for +shape: Shape.Shape
127  for tensor: Core.Tensor<shape>
128  for +coordinates: List<&2, Signed.Position>
129  {Signed.read_coordinate(shape, tensor, coordinates) == Core.read_coordinate(shape, tensor, Signed.wrap_each(Shape.dimensions(shape), coordinates)) :
130    Core.Tensor<shape> & Result<&2, &2, Core.Error, F32>}
131
132law signed_write_coordinate_wraps:
133  for +shape: Shape.Shape
134  for tensor: Core.Tensor<shape>
135  for +coordinates: List<&2, Signed.Position>
136  for value: F32
137  {Signed.write_coordinate(shape, tensor, coordinates, value) ==
138    Core.write_coordinate(shape, tensor, Signed.wrap_each(Shape.dimensions(shape), coordinates), value) :
139    Core.Tensor<shape> & Result<&2, &2, Core.Error, Unit>}
140
141law signed_read_flat_wraps:
142  for +shape: Shape.Shape
143  for tensor: Core.Tensor<shape>
144  for +index: Signed.Position
145  {Signed.read_flat(shape, tensor, index) == Core.read_flat(shape, tensor, Signed.wrap(Shape.count(shape), index)) :
146    Core.Tensor<shape> & Result<&2, &2, Core.Error, F32>}
147
148law signed_write_flat_wraps:
149  for +shape: Shape.Shape
150  for tensor: Core.Tensor<shape>
151  for +index: Signed.Position
152  for +value: F32
153  {Signed.write_flat(shape, tensor, index, value) == Core.write_flat(shape, tensor, Signed.wrap(Shape.count(shape), index), value) :
154    Core.Tensor<shape> & Result<&2, &2, Core.Error, Unit>}
155
156law signed_read_dynamic_wraps:
157  for input: Core.Dynamic
158  for +index: Signed.Position
159  {Signed.read_dynamic(input, index) == Core.read_dynamic(input, Signed.wrap(Shape.count(Totality.dynamic_shape(input)), index)) :
160    Core.Dynamic & Result<&2, &2, Core.Error, F32>}
161
162law signed_write_dynamic_wraps:
163  for input: Core.Dynamic
164  for +index: Signed.Position
165  for value: F32
166  {Signed.write_dynamic(input, index, value) == Core.write_dynamic(input, Signed.wrap(Shape.count(Totality.dynamic_shape(input)), index), value) :
167    Core.Dynamic & Result<&2, &2, Core.Error, Unit>}
168
169# A reshape infers its sizes from the input's element count, then is the
170# unsigned reshape.
171law signed_reshape_infers:
172  for +input: Shape.Shape
173  for tensor: Core.Tensor<input>
174  for +extents: List<&2, Signed.Extent>
175  {Signed.reshape(input, tensor, extents) ==
176    Signed.reshape_inferred(input, tensor, Signed.inferred_dimensions(Shape.elements(input), extents)) :
177    Result<&2, &1, Signed.Error, Core.Dynamic>}
178
179law signed_reshape_dynamic_infers:
180  for +input: Shape.Shape
181  for tensor: Core.Tensor<input>
182  for +extents: List<&2, Signed.Extent>
183  {Signed.reshape_dynamic(Core.Packed{input, tensor}, extents) == Signed.reshape(input, tensor, extents) :
184    Result<&2, &1, Signed.Error, Core.Dynamic>}
185
186# Exact sizes reshape exactly as the unsigned reshape does, its errors as
187# TransformError.
188law signed_reshape_unsigned:
189  for +input: Shape.Shape
190  for tensor: Core.Tensor<input>
191  for +dimensions: List<&2, U32>
192  {Signed.reshape(input, tensor, exacts(dimensions)) == Signed.transformed(Transforms.reshape(input, tensor, dimensions)) :
193    Result<&2, &1, Signed.Error, Core.Dynamic>}
194
195# With no Inferred size, the sizes are the request's.
196law signed_inferred_none:
197  for +count: Nat
198  for +extents: List<&2, Signed.Extent>
199  for none: {Signed.inferred_count(extents) == 0n : Nat}
200  {Signed.inferred_dimensions(count, extents) == Done{Signed.filled(extents, 0)} : Result<&2, &2, Signed.Error, List<&2, U32>>}
201
202# With one Inferred size and the other sizes multiplying to a product p > 0,
203# the Inferred size is the q with q * p = count.
204law signed_inferred_size:
205  for +count: Nat
206  for +extents: List<&2, Signed.Extent>
207  for one: {Signed.inferred_count(extents) == 1n : Nat}
208  for +rest: Nat
209  for product: {Signed.exact_product(extents) == 1n+rest : Nat}
210  for +quotient: Nat
211  for exact: {Nat.mul(quotient, 1n+rest) == count : Nat}
212  {Signed.inferred_dimensions(count, extents) == Done{Signed.filled(extents, U32.from_nat(quotient))} :
213    Result<&2, &2, Signed.Error, List<&2, U32>>}
214
215# And then the filled sizes hold q * p elements.
216law signed_filled_count:
217  for +extents: List<&2, Signed.Extent>
218  for one: {Signed.inferred_count(extents) == 1n : Nat}
219  for +size: U32
220  {Shape.logical_numel(Signed.filled(extents, size)) == Nat.mul(U32.to_nat(size), Signed.exact_product(extents)) : Nat}
221
222# p > 0 that does not divide count: ElementCountMismatch, as PyTorch's
223# "shape is invalid for input of size".
224law signed_inferred_indivisible:
225  for +count: Nat
226  for +extents: List<&2, Signed.Extent>
227  for one: {Signed.inferred_count(extents) == 1n : Nat}
228  for +rest: Nat
229  for product: {Signed.exact_product(extents) == 1n+rest : Nat}
230  for +quotient: Nat
231  for +remainder: Nat
232  for split: {Nat.divmod(count, 1n+rest) == (quotient, 1n+remainder) : Nat & Nat}
233  {Signed.inferred_dimensions(count, extents) == Fail{Signed.TransformError{Transforms.LayoutError{Layout.ElementCountMismatch{}}}} :
234    Result<&2, &2, Signed.Error, List<&2, U32>>}
235
236# p = 0: ambiguous on an empty input, ElementCountMismatch on any other.
237law signed_inferred_ambiguous:
238  for +extents: List<&2, Signed.Extent>
239  for one: {Signed.inferred_count(extents) == 1n : Nat}
240  for product: {Signed.exact_product(extents) == 0n : Nat}
241  {Signed.inferred_dimensions(0n, extents) == Fail{Signed.AmbiguousInferred{}} : Result<&2, &2, Signed.Error, List<&2, U32>>}
242
243law signed_inferred_zero_product:
244  for +count: Nat
245  for +extents: List<&2, Signed.Extent>
246  for one: {Signed.inferred_count(extents) == 1n : Nat}
247  for product: {Signed.exact_product(extents) == 0n : Nat}
248  {Signed.inferred_dimensions(1n+count, extents) == Fail{Signed.TransformError{Transforms.LayoutError{Layout.ElementCountMismatch{}}}} :
249    Result<&2, &2, Signed.Error, List<&2, U32>>}
250
251# The class of a product against a bound, in Signed.Product, stated from
252# its definition and not through signed.bend's helpers: 0 is ZeroProduct, a
253# nonzero product at most the bound is Within{product}, and one past it is
254# Beyond.
255def classified(product: Nat, below: Bool) -> Signed.Product:
256  match below:
257    case True{}:
258      Signed.Within{product}
259    case False{}:
260      Signed.Beyond{}
261
262def classify(+bound: Nat, product: Nat) -> Signed.Product:
263  match product:
264    case 0n:
265      Signed.ZeroProduct{}
266    case 1n+(+rest):
267      classified(1n+rest, Nat.is_le(1n+rest, bound))
268
269# The inference never forms the exact sizes' product: capped classifies it
270# against the input's element count, ZeroProduct, Within{p} with p at most
271# the count, or Beyond, multiplying only numbers at most the count. So the
272# inferred size is worked out from numbers at most 2^20, however large the
273# exact sizes, as [-1, 2^24, 2^24] of an empty input, whose product 2^48 is
274# past the native Nat limit.
275law signed_capped_classifies:
276  for +bound: Nat
277  for +extents: List<&2, Signed.Extent>
278  {Signed.capped(bound, extents) == classify(bound, Signed.exact_product(extents)) : Signed.Product}
279
280law signed_capped_within:
281  for +bound: Nat
282  for +extents: List<&2, Signed.Extent>
283  for +product: Nat
284  for within: {Signed.capped(bound, extents) == Signed.Within{product} : Signed.Product}
285  {Nat.is_le(product, bound) == True{} : Bool}
286
287# More than one Inferred size is refused before anything else.
288law signed_inferred_repeated:
289  for +count: Nat
290  for +extents: List<&2, Signed.Extent>
291  for +more: Nat
292  for repeated: {Signed.inferred_count(extents) == 2n+more : Nat}
293  {Signed.inferred_dimensions(count, extents) == Fail{Signed.RepeatedInferred{}} : Result<&2, &2, Signed.Error, List<&2, U32>>}
294
295# Filling puts the size at each Inferred place and keeps each exact size.
296law signed_filled_exact:
297  for +extent: U32
298  for +tail: List<&2, Signed.Extent>
299  for +size: U32
300  {Signed.filled(Signed.Exact{extent} <> tail, size) == extent <> Signed.filled(tail, size) : List<&2, U32>}
301
302law signed_filled_inferred:
303  for +tail: List<&2, Signed.Extent>
304  for +size: U32
305  {Signed.filled(Signed.Inferred{} <> tail, size) == size <> Signed.filled(tail, size) : List<&2, U32>}
306
307# The new error codes: RepeatedInferred 17, AmbiguousInferred 18; the
308# unsigned reshape's errors keep their codes.
309law signed_error_codes:
310  for error: Transforms.Error
311  {(Signed.error_code(Signed.TransformError{error}), Signed.error_code(Signed.RepeatedInferred{}), Signed.error_code(Signed.AmbiguousInferred{})) ==
312    (Transforms.error_code(error), 17, 18) : U32 & U32 & U32}
313
314# Signed slicing (Signed.slice, narrow and select). The axis is wrapped among
315# the input's rank, as every other signed axis; the positions along it among
316# its extent, each as PyTorch takes it.
317def axis_of(input: Core.Dynamic, dim: Signed.Position) -> U32:
318  Signed.wrap(Shape.rank(Totality.dynamic_shape(input)), dim)
319
320def extent_of(input: Core.Dynamic, +axis: U32) -> U32:
321  U32.from_nat(Slice.extent_at(Shape.dimensions(Totality.dynamic_shape(input)), axis))
322
323def rank_of(input: Core.Dynamic) -> Nat:
324  Shape.rank_dimensions(Shape.dimensions(Totality.dynamic_shape(input)))
325
326# A slice's start and stop, PyTorch's get_slice_bounds before its clamp to
327# the axis: a position counted from the start is kept; Last{k}, the position
328# d = -1 - k, becomes d + limit in range, and 0 below -limit.
329law signed_bound_first:
330  for +limit: U32
331  for +k: U32
332  {Signed.bound(limit, Signed.First{k}) == k : U32}
333
334law signed_bound_last_inside:
335  for +limit: U32
336  for +k: U32
337  for inside: {U32.is_lt(k, limit) == True{} : Bool}
338  {Nat.add(U32.to_nat(Signed.bound(limit, Signed.Last{k})), 1n+U32.to_nat(k)) == U32.to_nat(limit) : Nat}
339
340law signed_bound_last_outside:
341  for +limit: U32
342  for +k: U32
343  for outside: {U32.is_lt(k, limit) == False{} : Bool}
344  {Signed.bound(limit, Signed.Last{k}) == 0 : U32}
345
346# A signed slice is the unsigned slice at the wrapped axis of the bounds of
347# its start and stop among that axis's extent; slice_LAWS.bend's clamp then
348# gives PyTorch's.
349law signed_slice_bounds:
350  for input: Core.Dynamic
351  for +dim: Signed.Position
352  for start: Signed.Position
353  for stop: Signed.Position
354  for +step: U32
355  {Signed.slice(input, dim, start, stop, step) ==
356    Slice.slice(input, axis_of(input, dim), Signed.bound(extent_of(input, axis_of(input, dim)), start),
357      Signed.bound(extent_of(input, axis_of(input, dim)), stop), step) :
358    Core.Dynamic & Result<&2, &1, Slice.Error, Core.Dynamic>}
359
360# A signed select is the unsigned select at the wrapped axis of the index
361# wrapped among that axis's extent: an index outside -extent..extent-1 stays
362# at or past the extent, which the unsigned select refuses.
363law signed_select_wraps:
364  for input: Core.Dynamic
365  for +dim: Signed.Position
366  for index: Signed.Position
367  {Signed.select(input, dim, index) ==
368    Slice.select(input, axis_of(input, dim), Signed.wrap(extent_of(input, axis_of(input, dim)), index)) :
369    Core.Dynamic & Result<&2, &1, Slice.Error, Core.Dynamic>}
370
371# A signed narrow takes a start in -extent..extent, PyTorch's range. A start
372# counted from the start is the unsigned narrow's, which refuses one past
373# the extent itself; Last{k} below the extent is wrapped among it ...
374law signed_narrow_first:
375  for input: Core.Dynamic
376  for +dim: Signed.Position
377  for +k: U32
378  for +length: U32
379  {Signed.narrow(input, dim, Signed.First{k}, length) == Slice.narrow(input, axis_of(input, dim), k, length) :
380    Core.Dynamic & Result<&2, &1, Slice.Error, Core.Dynamic>}
381
382law signed_narrow_last_inside:
383  for input: Core.Dynamic
384  for +dim: Signed.Position
385  for +k: U32
386  for +length: U32
387  for inside: {U32.is_lt(k, extent_of(input, axis_of(input, dim))) == True{} : Bool}
388  {Signed.narrow(input, dim, Signed.Last{k}, length) ==
389    Slice.narrow(input, axis_of(input, dim), Signed.wrap(extent_of(input, axis_of(input, dim)), Signed.Last{k}), length) :
390    Core.Dynamic & Result<&2, &1, Slice.Error, Core.Dynamic>}
391
392# ... and a start before -extent is refused: OutOfRange once the axis is in
393# range, and AxisOutOfBounds first, as the unsigned narrow checks it.
394law signed_narrow_before:
395  for input: Core.Dynamic
396  for +dim: Signed.Position
397  for +k: U32
398  for +length: U32
399  for in_rank: {Nat.is_lt(U32.to_nat(axis_of(input, dim)), rank_of(input)) == True{} : Bool}
400  for before: {U32.is_lt(k, extent_of(input, axis_of(input, dim))) == False{} : Bool}
401  {Signed.narrow(input, dim, Signed.Last{k}, length) == (input, Fail{Slice.OutOfRange{}}) :
402    Core.Dynamic & Result<&2, &1, Slice.Error, Core.Dynamic>}
403
404law signed_narrow_before_axis:
405  for input: Core.Dynamic
406  for +dim: Signed.Position
407  for +k: U32
408  for +length: U32
409  for past_rank: {Nat.is_lt(U32.to_nat(axis_of(input, dim)), rank_of(input)) == False{} : Bool}
410  for before: {U32.is_lt(k, extent_of(input, axis_of(input, dim))) == False{} : Bool}
411  {Signed.narrow(input, dim, Signed.Last{k}, length) == (input, Fail{Slice.LayoutError{Layout.AxisOutOfBounds{}}}) :
412    Core.Dynamic & Result<&2, &1, Slice.Error, Core.Dynamic>}