src/tensor/slice_LAWS.bend

bendtorch-core@0.1.0.0 · 422 lines · 40 definitions

1import Base
2import bendtorch-core@0.1.0.0/src/shape/core.bend as Shape
3import bendtorch-core@0.1.0.0/src/storage/model.bend as Model
4import bendtorch-core@0.1.0.0/src/storage/refinement.bend as Refinement
5import bendtorch-core@0.1.0.0/src/tensor/core.bend as Core
6import bendtorch-core@0.1.0.0/src/tensor/layout.bend as Layout
7import bendtorch-core@0.1.0.0/src/tensor/transforms.bend as Transforms
8import bendtorch-core@0.1.0.0/src/tensor/elementwise_totality.bend as Totality
9import bendtorch-core@0.1.0.0/src/tensor/slice.bend as Slice
10
11# The laws of basic slicing (slice.bend), written from PyTorch's definition
12# of a slice: which start and stop it clamps to, how long the sliced axis
13# is, how the row-major view [outer, extent, inner] sits around the axis,
14# which calls it refuses, and what every cell of the forward copy and of the
15# backward scatter holds.
16
17type Holds<-P: Type> is Data:
18  Holds{-proof: P}
19
20def low(+clamped: Slice.Bounds) -> Nat:
21  Slice.Bounds{start, stop} = clamped
22  start
23
24def high(+clamped: Slice.Bounds) -> Nat:
25  Slice.Bounds{start, stop} = clamped
26  stop
27
28# PyTorch's clamp of a slice to an axis of extent n. A start at most n is
29# kept, a start past n becomes n.
30law slice_start_kept:
31  for +extent: Nat
32  for +start: Nat
33  for +stop: Nat
34  for inside: {Nat.is_le(start, extent) == True{} : Bool}
35  {low(Slice.bounds(extent, start, stop)) == start : Nat}
36
37law slice_start_clamped:
38  for +extent: Nat
39  for +start: Nat
40  for +stop: Nat
41  for past: {Nat.is_le(extent, start) == True{} : Bool}
42  {low(Slice.bounds(extent, start, stop)) == extent : Nat}
43
44# A stop from the start to n is kept, a stop past n becomes n, and a stop
45# at or before the start ends the slice at its start: an empty axis.
46law slice_stop_kept:
47  for +extent: Nat
48  for +start: Nat
49  for +stop: Nat
50  for +after: {Nat.is_le(start, stop) == True{} : Bool}
51  for +inside: {Nat.is_le(stop, extent) == True{} : Bool}
52  {high(Slice.bounds(extent, start, stop)) == stop : Nat}
53
54law slice_stop_clamped:
55  for +extent: Nat
56  for +start: Nat
57  for +stop: Nat
58  for past: {Nat.is_le(extent, stop) == True{} : Bool}
59  {high(Slice.bounds(extent, start, stop)) == extent : Nat}
60
61law slice_stop_empty:
62  for +extent: Nat
63  for +start: Nat
64  for +stop: Nat
65  for before: {Nat.is_le(stop, start) == True{} : Bool}
66  {high(Slice.bounds(extent, start, stop)) == low(Slice.bounds(extent, start, stop)) : Nat}
67
68# The length of the sliced axis is the number of indices start, start +
69# step, start + 2 step, ... below the stop: each index below the length
70# lies below the stop, and, for a positive step, the index at the length
71# does not.
72law slice_length_below:
73  for +start: Nat
74  for +stop: Nat
75  for +step: Nat
76  for +k: Nat
77  for below: {Nat.is_lt(k, Slice.length(Slice.Bounds{start, stop}, step)) == True{} : Bool}
78  {Nat.is_lt(Nat.add(start, Nat.mul(k, step)), stop) == True{} : Bool}
79
80law slice_length_reaches:
81  for +start: Nat
82  for +stop: Nat
83  for +step: Nat
84  for positive: {Nat.is_lt(0n, step) == True{} : Bool}
85  {Nat.is_le(stop, Nat.add(start, Nat.mul(Slice.length(Slice.Bounds{start, stop}, step), step))) == True{} : Bool}
86
87# The row-major view around axis `dim` of dimensions prefix ++ [extent] ++
88# suffix, with `dim` the rank of the prefix: the element count of the
89# prefix, the extent, and the element count of the suffix.
90law slice_view_around:
91  for +prefix: List<&2, U32>
92  for +extent: U32
93  for +suffix: List<&2, U32>
94  {Slice.view(List.append(&2, U32, prefix, extent <> suffix), Shape.rank_dimensions(prefix)) ==
95    Slice.View{Shape.logical_numel(prefix), U32.to_nat(extent), Shape.logical_numel(suffix)} : Slice.View}
96
97# The sliced output keeps every other axis and takes the new extent at the
98# axis; `select` drops the axis.
99law slice_resized_around:
100  for +prefix: List<&2, U32>
101  for +extent: U32
102  for +suffix: List<&2, U32>
103  for +other: U32
104  {Slice.resized(List.append(&2, U32, prefix, extent <> suffix), Shape.rank_dimensions(prefix), other) ==
105    List.append(&2, U32, prefix, other <> suffix) : List<&2, U32>}
106
107law select_dropped_around:
108  for +prefix: List<&2, U32>
109  for +extent: U32
110  for +suffix: List<&2, U32>
111  {Slice.dropped(List.append(&2, U32, prefix, extent <> suffix), Shape.rank_dimensions(prefix)) ==
112    List.append(&2, U32, prefix, suffix) : List<&2, U32>}
113
114# What a plan holds along its axis (the extent, the clamped start, the step
115# and the length) and around it (the element counts before and after it).
116def axis_part(+plan: Slice.Plan) -> Nat & Nat & Nat & Nat:
117  Slice.Plan{outer, extent, inner, start, step, length} = plan
118  (extent, start, step, length)
119
120def view_part(+plan: Slice.Plan) -> Nat & Nat:
121  Slice.Plan{outer, extent, inner, start, step, length} = plan
122  (outer, inner)
123
124# PyTorch's clamp of a slice position to an axis of extent n: the position
125# if it is at most n, else n.
126def clamp(+position: Nat, +extent: Nat) -> Nat:
127  Bool.pick(Nat, Nat.is_le(position, extent), position, extent)
128
129# The stop after the clamp: the clamped stop, or the clamped start where the
130# clamped stop falls before it, so the axis is empty.
131def clamped_stop(+start: Nat, +stop: Nat, +extent: Nat) -> Nat:
132  Bool.pick(Nat, Nat.is_le(clamp(start, extent), clamp(stop, extent)), clamp(stop, extent), clamp(start, extent))
133
134# What a request start:stop:step plans along an axis of extent n: the
135# extent, the clamped start, the step, and the length of the clamped range.
136def axis_spec(+extent: Nat, +start: Nat, +stop: Nat, +step: Nat) -> Nat & Nat & Nat & Nat:
137  (extent, clamp(start, extent), step,
138    Slice.length(Slice.Bounds{clamp(start, extent), clamped_stop(start, stop, extent)}, step))
139
140# The plan of a request at axis `dim` of prefix ++ [extent] ++ suffix, with
141# `dim` the rank of the prefix, holds the request's own start, stop and step,
142# clamped to that extent, whatever the other extents, an empty input's
143# included.
144law slice_plan_axis:
145  for +prefix: List<&2, U32>
146  for +extent: U32
147  for +suffix: List<&2, U32>
148  for +dim: U32
149  for at: {U32.to_nat(dim) == Shape.rank_dimensions(prefix) : Nat}
150  for +start: U32
151  for +stop: U32
152  for +step: U32
153  {axis_part(Slice.plan_of(List.append(&2, U32, prefix, extent <> suffix), dim, start, stop, step)) ==
154    axis_spec(U32.to_nat(extent), U32.to_nat(start), U32.to_nat(stop), U32.to_nat(step)) : Nat & Nat & Nat & Nat}
155
156# For an input with an element, the plan's view is the element counts of the
157# prefix and of the suffix, the row-major view around the axis.
158law slice_plan_view:
159  for +prefix: List<&2, U32>
160  for +extent: U32
161  for +suffix: List<&2, U32>
162  for +dim: U32
163  for at: {U32.to_nat(dim) == Shape.rank_dimensions(prefix) : Nat}
164  for +start: U32
165  for +stop: U32
166  for +step: U32
167  for nonempty: {Nat.is_lt(0n, Shape.logical_numel(List.append(&2, U32, prefix, extent <> suffix))) == True{} : Bool}
168  {view_part(Slice.plan_of(List.append(&2, U32, prefix, extent <> suffix), dim, start, stop, step)) ==
169    (Shape.logical_numel(prefix), Shape.logical_numel(suffix)) : Nat & Nat}
170
171# An input with no element gets a plan with no view: neither build has a
172# cell then, so the plan forms no product of the other extents, which can
173# pass the largest Nat the runtime holds although no element exists.
174law slice_plan_empty_view:
175  for +dimensions: List<&2, U32>
176  for +dim: U32
177  for +start: U32
178  for +stop: U32
179  for +step: U32
180  for empty: {Shape.logical_numel(dimensions) == 0n : Nat}
181  {view_part(Slice.plan_of(dimensions, dim, start, stop, step)) == (0n, 0n) : Nat & Nat}
182
183# PyTorch's checks, in its order. An axis at or past the rank (every axis
184# of a rank-0 input) is refused first, then a step of 0.
185law slice_axis_refused:
186  for +input: Shape.Shape
187  for +dim: U32
188  for +start: U32
189  for +stop: U32
190  for +step: U32
191  for outside: {Nat.is_lt(U32.to_nat(dim), Shape.rank_dimensions(Shape.dimensions(input))) == False{} : Bool}
192  {Slice.plan(input, dim, start, stop, step) == Fail{Slice.LayoutError{Layout.AxisOutOfBounds{}}} :
193    Result<&2, &2, Slice.Error, Slice.Planned>}
194
195law slice_step_refused:
196  for +input: Shape.Shape
197  for +dim: U32
198  for +start: U32
199  for +stop: U32
200  for +step: U32
201  for inside: {Nat.is_lt(U32.to_nat(dim), Shape.rank_dimensions(Shape.dimensions(input))) == True{} : Bool}
202  for zero: {U32.to_nat(step) == 0n : Nat}
203  {Slice.plan(input, dim, start, stop, step) == Fail{Slice.StepNotPositive{}} : Result<&2, &2, Slice.Error, Slice.Planned>}
204
205# What an accepted slice plans: the view around the axis, the start and
206# stop clamped to the extent, the step, the length of the sliced axis, and
207# the output's dimensions, the input's with that length at the axis.
208def PlannedAs(+dimensions: List<&2, U32>, +dim: U32, +start: U32, +stop: U32, +step: U32,
209  result: Result<&2, &2, Slice.Error, Slice.Planned>) -> Type:
210  match result:
211    case Fail{error}:
212      Unit
213    case Done{Slice.Planned{plan, output}}:
214      {plan == Slice.plan_of(dimensions, dim, start, stop, step) : Slice.Plan} &
215        {Shape.dimensions(output) == Slice.resized(dimensions, U32.to_nat(dim), U32.from_nat(Slice.plan_length(plan))) : List<&2, U32>}
216
217law slice_planned:
218  for +input: Shape.Shape
219  for +dim: U32
220  for +start: U32
221  for +stop: U32
222  for +step: U32
223  PlannedAs(Shape.dimensions(input), dim, start, stop, step, Slice.plan(input, dim, start, stop, step))
224
225# A refused plan fails the slice with its error and hands the input back.
226law slice_refused:
227  for input: Core.Dynamic
228  for +dim: U32
229  for +start: U32
230  for +stop: U32
231  for +step: U32
232  for error: Slice.Error
233  for refused: {Slice.plan(Totality.dynamic_shape(input), dim, start, stop, step) == Fail{error} : Result<&2, &2, Slice.Error, Slice.Planned>}
234  {Slice.slice(input, dim, start, stop, step) == (input, Fail{error}) : Core.Dynamic & Result<&2, &1, Slice.Error, Core.Dynamic>}
235
236def rank_of(input: Core.Dynamic) -> Nat:
237  Shape.rank_dimensions(Shape.dimensions(Totality.dynamic_shape(input)))
238
239def extent_of(input: Core.Dynamic, +dim: U32) -> Nat:
240  Slice.extent_at(Shape.dimensions(Totality.dynamic_shape(input)), dim)
241
242# PyTorch's narrow(dim, start, length) is the slice start:start + length
243# with step 1, once the axis is in range and start + length is at most the
244# extent; a start + length past the extent is refused with OutOfRange, an
245# axis past the rank as a slice refuses it.
246law narrow_slices:
247  for input: Core.Dynamic
248  for +dim: U32
249  for +start: U32
250  for +length: U32
251  for inside: {Nat.is_lt(U32.to_nat(dim), rank_of(input)) == True{} : Bool}
252  for fits: {Nat.is_le(Nat.add(U32.to_nat(start), U32.to_nat(length)), extent_of(input, dim)) == True{} : Bool}
253  {Slice.narrow(input, dim, start, length) == Slice.slice(input, dim, start, U32.add(start, length), 1) :
254    Core.Dynamic & Result<&2, &1, Slice.Error, Core.Dynamic>}
255
256law narrow_out_of_range:
257  for input: Core.Dynamic
258  for +dim: U32
259  for +start: U32
260  for +length: U32
261  for inside: {Nat.is_lt(U32.to_nat(dim), rank_of(input)) == True{} : Bool}
262  for past: {Nat.is_le(Nat.add(U32.to_nat(start), U32.to_nat(length)), extent_of(input, dim)) == False{} : Bool}
263  {Slice.narrow(input, dim, start, length) == (input, Fail{Slice.OutOfRange{}}) : Core.Dynamic & Result<&2, &1, Slice.Error, Core.Dynamic>}
264
265law narrow_axis_refused:
266  for input: Core.Dynamic
267  for +dim: U32
268  for +start: U32
269  for +length: U32
270  for outside: {Nat.is_lt(U32.to_nat(dim), rank_of(input)) == False{} : Bool}
271  {Slice.narrow(input, dim, start, length) == (input, Fail{Slice.LayoutError{Layout.AxisOutOfBounds{}}}) :
272    Core.Dynamic & Result<&2, &1, Slice.Error, Core.Dynamic>}
273
274# A slice's result with its output reshaped to `dimensions`.
275def reshape_as_slice(result: Result<&2, &1, Transforms.Error, Core.Dynamic>) -> Result<&2, &1, Slice.Error, Core.Dynamic>:
276  match result:
277    case Fail{Transforms.LayoutError{error}}:
278      Fail{Slice.LayoutError{error}}
279    case Fail{Transforms.TensorError{error}}:
280      Fail{Slice.TensorError{error}}
281    case Done{output}:
282      Done{output}
283
284def reshaped(+dimensions: List<&2, U32>, pair: Core.Dynamic & Result<&2, &1, Slice.Error, Core.Dynamic>) ->
285  Core.Dynamic & Result<&2, &1, Slice.Error, Core.Dynamic>:
286  (input, result) = pair
287  match result:
288    case Fail{error}:
289      (input, Fail{error})
290    case Done{output}:
291      (input, reshape_as_slice(Transforms.reshape_dynamic(output, dimensions)))
292
293# PyTorch's select(dim, index) is the slice index:index + 1 with step 1,
294# reshaped to the input's dimensions without the axis, once the axis is in
295# range and the index below the extent; an index at or past the extent is
296# refused with OutOfRange, an axis past the rank as a slice refuses it.
297law select_slices:
298  for input: Core.Dynamic
299  for +dim: U32
300  for +index: U32
301  for inside: {Nat.is_lt(U32.to_nat(dim), rank_of(input)) == True{} : Bool}
302  for below: {Nat.is_lt(U32.to_nat(index), extent_of(input, dim)) == True{} : Bool}
303  {Slice.select(input, dim, index) ==
304    reshaped(Slice.dropped(Shape.dimensions(Totality.dynamic_shape(input)), U32.to_nat(dim)), Slice.slice(input, dim, index, U32.add(index, 1), 1)) :
305    Core.Dynamic & Result<&2, &1, Slice.Error, Core.Dynamic>}
306
307law select_out_of_range:
308  for input: Core.Dynamic
309  for +dim: U32
310  for +index: U32
311  for inside: {Nat.is_lt(U32.to_nat(dim), rank_of(input)) == True{} : Bool}
312  for past: {Nat.is_lt(U32.to_nat(index), extent_of(input, dim)) == False{} : Bool}
313  {Slice.select(input, dim, index) == (input, Fail{Slice.OutOfRange{}}) : Core.Dynamic & Result<&2, &1, Slice.Error, Core.Dynamic>}
314
315law select_axis_refused:
316  for input: Core.Dynamic
317  for +dim: U32
318  for +index: U32
319  for outside: {Nat.is_lt(U32.to_nat(dim), rank_of(input)) == False{} : Bool}
320  {Slice.select(input, dim, index) == (input, Fail{Slice.LayoutError{Layout.AxisOutOfBounds{}}}) :
321    Core.Dynamic & Result<&2, &1, Slice.Error, Core.Dynamic>}
322
323# The cells of the builds, in the coordinates of the views: output cell
324# (o, k, r) of [outer, length, inner] sits at row-major position
325# (o * length + k) * inner + r, and input cell (o, i, r) of [outer, extent,
326# inner] at (o * extent + i) * inner + r.
327def position(+rows: Nat, +inner: Nat, +o: Nat, +i: Nat, +r: Nat) -> Nat:
328  Nat.add(Nat.mul(Nat.add(Nat.mul(o, rows), i), inner), r)
329
330# For a balanced input array A and every split, the slot of the forward
331# build that holds output cell (o, k, r), below the count, holds input cell
332# (o, start + k * step, r): the word PyTorch's slice puts there.
333law slice_build_cells:
334  for +depth: Nat
335  for +splits: Nat
336  for +outer: Nat
337  for +extent: Nat
338  for +inner: Nat
339  for +start: Nat
340  for +step: Nat
341  for +length: Nat
342  for +count: Nat
343  for +base: Nat
344  for +da: Nat
345  for -A: Array<F32>
346  for a: Array<F32>
347  for same: {a == A : Array<F32>}
348  for valid: Model.Balanced(F32, da, A)
349  for +path: Model.Path(depth)
350  for +o: Nat
351  for +k: Nat
352  for +r: Nat
353  for k_below: {Nat.is_lt(k, length) == True{} : Bool}
354  for r_below: {Nat.is_lt(r, inner) == True{} : Bool}
355  for at: {Nat.add(Model.offset(depth, path), base) == position(length, inner, o, k, r) : Nat}
356  for inside: {Nat.is_lt(Nat.add(Model.offset(depth, path), base), count) == True{} : Bool}
357  Holds<{Model.observe(F32, depth, Slice.build(depth, splits, Slice.Plan{outer, extent, inner, start, step, length}, count, base, a), path) ==
358    Model.observe(F32, da, A, Refinement.route(da, U32.from_nat(position(extent, inner, o, Nat.add(start, Nat.mul(k, step)), r)))) :
359    Maybe<&2, F32>}>
360
361# For a balanced gradient array G and every split, the slot of the backward
362# build that holds input cell (o, start + k * step, r), with k below the
363# length and a positive step, holds gradient cell (o, k, r) ...
364law slice_scatter_hit:
365  for +depth: Nat
366  for +splits: Nat
367  for +outer: Nat
368  for +extent: Nat
369  for +inner: Nat
370  for +start: Nat
371  for +step: Nat
372  for +length: Nat
373  for +count: Nat
374  for +base: Nat
375  for +dg: Nat
376  for -G: Array<F32>
377  for g: Array<F32>
378  for same: {g == G : Array<F32>}
379  for valid: Model.Balanced(F32, dg, G)
380  for +path: Model.Path(depth)
381  for +o: Nat
382  for +k: Nat
383  for +r: Nat
384  for positive: {Nat.is_lt(0n, step) == True{} : Bool}
385  for k_below: {Nat.is_lt(k, length) == True{} : Bool}
386  for i_below: {Nat.is_lt(Nat.add(start, Nat.mul(k, step)), extent) == True{} : Bool}
387  for r_below: {Nat.is_lt(r, inner) == True{} : Bool}
388  for at: {Nat.add(Model.offset(depth, path), base) == position(extent, inner, o, Nat.add(start, Nat.mul(k, step)), r) : Nat}
389  for inside: {Nat.is_lt(Nat.add(Model.offset(depth, path), base), count) == True{} : Bool}
390  Holds<{Model.observe(F32, depth, Slice.scatter_build(depth, splits, Slice.Plan{outer, extent, inner, start, step, length}, count, base, g), path) ==
391    Model.observe(F32, dg, G, Refinement.route(dg, U32.from_nat(position(length, inner, o, k, r)))) : Maybe<&2, F32>}>
392
393# ... and the slot of input cell (o, i, r) where no k below the length has
394# i = start + k * step holds +0.0, as PyTorch's slice_backward leaves it.
395law slice_scatter_miss:
396  for +depth: Nat
397  for +splits: Nat
398  for +outer: Nat
399  for +extent: Nat
400  for +inner: Nat
401  for +start: Nat
402  for +step: Nat
403  for +length: Nat
404  for +count: Nat
405  for +base: Nat
406  for +dg: Nat
407  for -G: Array<F32>
408  for g: Array<F32>
409  for same: {g == G : Array<F32>}
410  for valid: Model.Balanced(F32, dg, G)
411  for +path: Model.Path(depth)
412  for +o: Nat
413  for +i: Nat
414  for +r: Nat
415  for i_below: {Nat.is_lt(i, extent) == True{} : Bool}
416  for r_below: {Nat.is_lt(r, inner) == True{} : Bool}
417  for missed: @+k: Nat -> @below: {Nat.is_lt(k, length) == True{} : Bool} ->
418    {Nat.is_eq(Nat.add(start, Nat.mul(k, step)), i) == False{} : Bool}
419  for at: {Nat.add(Model.offset(depth, path), base) == position(extent, inner, o, i, r) : Nat}
420  for inside: {Nat.is_lt(Nat.add(Model.offset(depth, path), base), count) == True{} : Bool}
421  Holds<{Model.observe(F32, depth, Slice.scatter_build(depth, splits, Slice.Plan{outer, extent, inner, start, step, length}, count, base, g), path) ==
422    Some{0.0} : Maybe<&2, F32>}>