src/tensor/gather_LAWS.bend

bendtorch-core@0.1.0.0 · 339 lines · 31 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/elementwise_totality.bend as Totality
8import bendtorch-core@0.1.0.0/src/tensor/slice.bend as Slice
9import bendtorch-core@0.1.0.0/src/tensor/signed.bend as Signed
10import bendtorch-core@0.1.0.0/src/tensor/gather.bend as Gather
11
12# The laws of index_select and gather (gather.bend), written from PyTorch's
13# definitions: which requests they refuse and in which order, the plan and
14# the output shape of the ones they take, what every cell of the forward
15# copy holds, and what every cell of the backward, index_add and
16# scatter_add into +0.0, holds.
17
18type Holds<-P: Type> is Data:
19  Holds{-proof: P}
20
21# The axis a request names: its signed dim wrapped among the input's rank.
22def axis(+input: Shape.Shape, dim: Signed.Position) -> Nat:
23  U32.to_nat(Signed.wrap(Shape.rank(input), dim))
24
25def rank(+input: Shape.Shape) -> Nat:
26  Shape.rank_dimensions(Shape.dimensions(input))
27
28def extent(+input: Shape.Shape, +at: Nat) -> Nat:
29  Slice.axis_extent(Shape.dimensions(input), at)
30
31# How many entries an index holds.
32def length(entries: List<&2, U32>) -> Nat:
33  match entries:
34    case Nil{}:
35      0n
36    case head <> rest:
37      1n+length(rest)
38
39# Whether every entry lies below the extent.
40def inside(entries: List<&2, U32>, +bound: Nat) -> Bool:
41  match entries:
42    case Nil{}:
43      True{}
44    case head <> rest:
45      Bool.and(Nat.is_lt(U32.to_nat(head), bound), inside(rest, bound))
46
47# Whether two lists of extents are equal, extent by extent.
48def equal(left: List<&2, U32>, right: List<&2, U32>) -> Bool:
49  match left right:
50    case Nil{} Nil{}:
51      True{}
52    case x <> xs y <> ys:
53      Bool.and(Nat.is_eq(U32.to_nat(x), U32.to_nat(y)), equal(xs, ys))
54    case _ _:
55      False{}
56
57# Whether a gather index has the input's extent at every axis but `at`.
58def others(dimensions: List<&2, U32>, index: List<&2, U32>, +at: Nat) -> Bool:
59  match dimensions index at:
60    case x <> xs y <> ys 0n:
61      equal(xs, ys)
62    case x <> xs y <> ys 1n+next:
63      Bool.and(Nat.is_eq(U32.to_nat(x), U32.to_nat(y)), others(xs, ys, next))
64    case _ _ _:
65      False{}
66
67# The plan of a request the checks take: the row-major view of the input
68# around the axis, slice.bend's, and the length of the output's axis ...
69def viewed(+kind: Gather.Lookup, +given: Slice.View, +count: Nat) -> Gather.Plan:
70  Slice.View{+outer, +rows, +inner} = given
71  Gather.Plan{kind, outer, rows, inner, count}
72
73# ... and its output, the input's dimensions with that length at the axis,
74# or the refusal of the shape's own checks (rank and capacity).
75def output(-dimensions: List<&2, U32>, +plan: Gather.Plan, built: Shape.ConstructResult<dimensions>) ->
76  Result<&2, &2, Gather.Error, Gather.Planned>:
77  match built:
78    case Shape.SError{error}:
79      Fail{Gather.TensorError{Core.ShapeError{error}}}
80    case Shape.SBuilt{+shape, preserved}:
81      Done{Gather.Planned{plan, shape}}
82
83def taken(+kind: Gather.Lookup, +dimensions: List<&2, U32>, +at: Nat, +count: Nat) -> Result<&2, &2, Gather.Error, Gather.Planned>:
84  output(Slice.resized(dimensions, at, U32.from_nat(count)),
85    viewed(kind, Slice.sized(dimensions, at, Slice.holds_zero(dimensions)), count),
86    Shape.construct(Slice.resized(dimensions, at, U32.from_nat(count))))
87
88# index_select(dim, index): an axis outside -rank..rank-1 is refused first,
89# then an entry at or past the extent; any other request is planned.
90law index_select_axis_refused:
91  for +input: Shape.Shape
92  for +dim: Signed.Position
93  for +entries: List<&2, U32>
94  for outside: {Nat.is_lt(axis(input, dim), rank(input)) == False{} : Bool}
95  {Gather.select_plan(input, dim, entries) == Fail{Gather.LayoutError{Layout.AxisOutOfBounds{}}} :
96    Result<&2, &2, Gather.Error, Gather.Planned>}
97
98law index_select_out_of_range:
99  for +input: Shape.Shape
100  for +dim: Signed.Position
101  for +entries: List<&2, U32>
102  for in_rank: {Nat.is_lt(axis(input, dim), rank(input)) == True{} : Bool}
103  for past: {inside(entries, extent(input, axis(input, dim))) == False{} : Bool}
104  {Gather.select_plan(input, dim, entries) == Fail{Gather.OutOfRange{}} : Result<&2, &2, Gather.Error, Gather.Planned>}
105
106law index_select_planned:
107  for +input: Shape.Shape
108  for +dim: Signed.Position
109  for +entries: List<&2, U32>
110  for in_rank: {Nat.is_lt(axis(input, dim), rank(input)) == True{} : Bool}
111  for within: {inside(entries, extent(input, axis(input, dim))) == True{} : Bool}
112  {Gather.select_plan(input, dim, entries) == taken(Gather.Selected{}, Shape.dimensions(input), axis(input, dim), length(entries)) :
113    Result<&2, &2, Gather.Error, Gather.Planned>}
114
115# gather(dim, index): an entry count other than the element count of the
116# index's shape is refused first, then the axis, an index of another rank, one
117# with another extent at another axis, and an entry at or past the extent;
118# any other request is planned, with the index's extent at the axis.
119law gather_count_refused:
120  for +input: Shape.Shape
121  for +dim: Signed.Position
122  for +index: Shape.Shape
123  for +entries: List<&2, U32>
124  for miscounted: {Nat.is_eq(length(entries), Shape.elements(index)) == False{} : Bool}
125  {Gather.gather_plan(input, dim, Gather.Index{index, entries}) == Fail{Gather.TensorError{Core.ValueCountMismatch{}}} :
126    Result<&2, &2, Gather.Error, Gather.Planned>}
127
128law gather_axis_refused:
129  for +input: Shape.Shape
130  for +dim: Signed.Position
131  for +index: Shape.Shape
132  for +entries: List<&2, U32>
133  for counted: {Nat.is_eq(length(entries), Shape.elements(index)) == True{} : Bool}
134  for outside: {Nat.is_lt(axis(input, dim), rank(input)) == False{} : Bool}
135  {Gather.gather_plan(input, dim, Gather.Index{index, entries}) == Fail{Gather.LayoutError{Layout.AxisOutOfBounds{}}} :
136    Result<&2, &2, Gather.Error, Gather.Planned>}
137
138law gather_rank_refused:
139  for +input: Shape.Shape
140  for +dim: Signed.Position
141  for +index: Shape.Shape
142  for +entries: List<&2, U32>
143  for counted: {Nat.is_eq(length(entries), Shape.elements(index)) == True{} : Bool}
144  for in_rank: {Nat.is_lt(axis(input, dim), rank(input)) == True{} : Bool}
145  for ranked: {Nat.is_eq(Shape.rank_dimensions(Shape.dimensions(index)), rank(input)) == False{} : Bool}
146  {Gather.gather_plan(input, dim, Gather.Index{index, entries}) == Fail{Gather.IndexRankMismatch{}} :
147    Result<&2, &2, Gather.Error, Gather.Planned>}
148
149law gather_shape_refused:
150  for +input: Shape.Shape
151  for +dim: Signed.Position
152  for +index: Shape.Shape
153  for +entries: List<&2, U32>
154  for counted: {Nat.is_eq(length(entries), Shape.elements(index)) == True{} : Bool}
155  for in_rank: {Nat.is_lt(axis(input, dim), rank(input)) == True{} : Bool}
156  for ranked: {Nat.is_eq(Shape.rank_dimensions(Shape.dimensions(index)), rank(input)) == True{} : Bool}
157  for unmatched: {others(Shape.dimensions(input), Shape.dimensions(index), axis(input, dim)) == False{} : Bool}
158  {Gather.gather_plan(input, dim, Gather.Index{index, entries}) == Fail{Gather.IndexShapeMismatch{}} :
159    Result<&2, &2, Gather.Error, Gather.Planned>}
160
161law gather_out_of_range:
162  for +input: Shape.Shape
163  for +dim: Signed.Position
164  for +index: Shape.Shape
165  for +entries: List<&2, U32>
166  for counted: {Nat.is_eq(length(entries), Shape.elements(index)) == True{} : Bool}
167  for in_rank: {Nat.is_lt(axis(input, dim), rank(input)) == True{} : Bool}
168  for ranked: {Nat.is_eq(Shape.rank_dimensions(Shape.dimensions(index)), rank(input)) == True{} : Bool}
169  for matched: {others(Shape.dimensions(input), Shape.dimensions(index), axis(input, dim)) == True{} : Bool}
170  for past: {inside(entries, extent(input, axis(input, dim))) == False{} : Bool}
171  {Gather.gather_plan(input, dim, Gather.Index{index, entries}) == Fail{Gather.OutOfRange{}} :
172    Result<&2, &2, Gather.Error, Gather.Planned>}
173
174law gather_planned:
175  for +input: Shape.Shape
176  for +dim: Signed.Position
177  for +index: Shape.Shape
178  for +entries: List<&2, U32>
179  for counted: {Nat.is_eq(length(entries), Shape.elements(index)) == True{} : Bool}
180  for in_rank: {Nat.is_lt(axis(input, dim), rank(input)) == True{} : Bool}
181  for ranked: {Nat.is_eq(Shape.rank_dimensions(Shape.dimensions(index)), rank(input)) == True{} : Bool}
182  for matched: {others(Shape.dimensions(input), Shape.dimensions(index), axis(input, dim)) == True{} : Bool}
183  for within: {inside(entries, extent(input, axis(input, dim))) == True{} : Bool}
184  {Gather.gather_plan(input, dim, Gather.Index{index, entries}) ==
185    taken(Gather.Gathered{}, Shape.dimensions(input), axis(input, dim), Slice.axis_extent(Shape.dimensions(index), axis(input, dim))) :
186    Result<&2, &2, Gather.Error, Gather.Planned>}
187
188# A gather index is built as a tensor of its dimensions is: refused past
189# rank 4 or 2^20 elements with the shape's own error, else it holds the
190# shape of those dimensions.
191def index_of(+entries: List<&2, U32>, -dimensions: List<&2, U32>, built: Shape.ConstructResult<dimensions>) ->
192  Result<&2, &2, Gather.Error, Gather.Index>:
193  match built:
194    case Shape.SError{error}:
195      Fail{Gather.TensorError{Core.ShapeError{error}}}
196    case Shape.SBuilt{+shape, preserved}:
197      Done{Gather.Index{shape, entries}}
198
199law gather_index_shaped:
200  for +dimensions: List<&2, U32>
201  for +entries: List<&2, U32>
202  {Gather.index(dimensions, entries) == index_of(entries, dimensions, Shape.construct(dimensions)) : Result<&2, &2, Gather.Error, Gather.Index>}
203
204# A refused plan fails the copy with its error and hands the input back.
205law index_select_refused:
206  for input: Core.Dynamic
207  for +dim: Signed.Position
208  for +entries: List<&2, U32>
209  for error: Gather.Error
210  for refused: {Gather.select_plan(Totality.dynamic_shape(input), dim, entries) == Fail{error} : Result<&2, &2, Gather.Error, Gather.Planned>}
211  {Gather.index_select(input, dim, entries) == (input, Fail{error}) : Core.Dynamic & Result<&2, &1, Gather.Error, Core.Dynamic>}
212
213law gather_refused:
214  for input: Core.Dynamic
215  for +dim: Signed.Position
216  for +index: Shape.Shape
217  for +entries: List<&2, U32>
218  for error: Gather.Error
219  for refused: {Gather.gather_plan(Totality.dynamic_shape(input), dim, Gather.Index{index, entries}) == Fail{error} :
220    Result<&2, &2, Gather.Error, Gather.Planned>}
221  {Gather.gather(input, dim, Gather.Index{index, entries}) == (input, Fail{error}) : Core.Dynamic & Result<&2, &1, Gather.Error, Core.Dynamic>}
222
223# The cells of the builds, in the coordinates of the views: cell (o, k, r)
224# of [outer, rows, inner] sits at (o * rows + k) * inner + r.
225def position(+rows: Nat, +inner: Nat, +o: Nat, +k: Nat, +r: Nat) -> Nat:
226  Nat.add(Nat.mul(Nat.add(Nat.mul(o, rows), k), inner), r)
227
228# The entry output cell (o, k, r) of [outer, length, inner] reads: entry k
229# of an index_select's index, the entry at the cell itself of a gather's.
230def entry_at(+kind: Gather.Lookup, +count: Nat, +inner: Nat, +o: Nat, +k: Nat, +r: Nat) -> Nat:
231  match kind:
232    case Gather.Selected{}:
233      k
234    case Gather.Gathered{}:
235      position(count, inner, o, k, r)
236
237# For a balanced input array A, a balanced index array I and every split,
238# the slot of the forward build that holds output cell (o, k, r), below
239# the count, holds input cell (o, j, r), where j is the entry the cell
240# reads: the word PyTorch's index_select and gather put there.
241law gather_build_cells:
242  for +depth: Nat
243  for +splits: Nat
244  for +kind: Gather.Lookup
245  for +outer: Nat
246  for +rows: Nat
247  for +inner: Nat
248  for +count: Nat
249  for +cells: Nat
250  for +base: Nat
251  for +da: Nat
252  for -A: Array<F32>
253  for a: Array<F32>
254  for same: {a == A : Array<F32>}
255  for valid: Model.Balanced(F32, da, A)
256  for +di: Nat
257  for -I: Array<U32>
258  for index: Array<U32>
259  for same_index: {index == I : Array<U32>}
260  for index_valid: Model.Balanced(U32, di, I)
261  for +path: Model.Path(depth)
262  for +o: Nat
263  for +k: Nat
264  for +r: Nat
265  for k_below: {Nat.is_lt(k, count) == True{} : Bool}
266  for r_below: {Nat.is_lt(r, inner) == True{} : Bool}
267  for at: {Nat.add(Model.offset(depth, path), base) == position(count, inner, o, k, r) : Nat}
268  for in_count: {Nat.is_lt(Nat.add(Model.offset(depth, path), base), cells) == True{} : Bool}
269  for +j: U32
270  for read: {Model.observe(U32, di, I, Refinement.route(di, U32.from_nat(entry_at(kind, count, inner, o, k, r)))) == Some{j} : Maybe<&2, U32>}
271  Holds<{Model.observe(F32, depth, Gather.build(depth, splits, Gather.Plan{kind, outer, rows, inner, count}, cells, base, a, index), path) ==
272    Model.observe(F32, da, A, Refinement.route(da, U32.from_nat(position(rows, inner, o, U32.to_nat(j), r)))) : Maybe<&2, F32>}>
273
274# The backward of input cell (o, i, r): from gradient cell (o, k, r) on, for
275# `left` more k, a cell whose entry is i adds its word to the sum, in
276# increasing k; a cell whose entry is another skips. `value` is the sum at
277# the end.
278def Added(+dg: Nat, -G: Array<F32>, +di: Nat, -I: Array<U32>, +kind: Gather.Lookup, +count: Nat, +inner: Nat, +o: Nat, +i: Nat, +r: Nat,
279  left: Nat, +k: Nat, +acc: F32, +value: F32) -> Type:
280  match left:
281    case 0n:
282      {value == acc : F32}
283    case 1n+rest:
284      AddStep<dg, G, di, I, kind, count, inner, o, i, r, rest, k, acc, value>
285
286type AddStep<-dg: Nat, -G: Array<F32>, -di: Nat, -I: Array<U32>, -kind: Gather.Lookup, -count: Nat, -inner: Nat, -o: Nat, -i: Nat, -r: Nat,
287  -rest: Nat, -k: Nat, -acc: F32, -value: F32> is Data:
288  Skips{-j: U32,
289    -read: {Model.observe(U32, di, I, Refinement.route(di, U32.from_nat(entry_at(kind, count, inner, o, k, r)))) == Some{j} : Maybe<&2, U32>},
290    -other: {Nat.is_eq(U32.to_nat(j), i) == False{} : Bool},
291    -next: Added(dg, G, di, I, kind, count, inner, o, i, r, rest, 1n+k, acc, value)}
292  Adds{-j: U32,
293    -read: {Model.observe(U32, di, I, Refinement.route(di, U32.from_nat(entry_at(kind, count, inner, o, k, r)))) == Some{j} : Maybe<&2, U32>},
294    -same: {Nat.is_eq(U32.to_nat(j), i) == True{} : Bool},
295    -x: F32,
296    -word: {Model.observe(F32, dg, G, Refinement.route(dg, U32.from_nat(position(count, inner, o, k, r)))) == Some{x} : Maybe<&2, F32>},
297    -next: Added(dg, G, di, I, kind, count, inner, o, i, r, rest, 1n+k, F32.add(acc, x), value)}
298
299# A slot that holds `value`, the sum Added names from +0.0 over every k.
300type Summed<-depth: Nat, -out: Array<F32>, -path: Model.Path(depth), -dg: Nat, -G: Array<F32>, -di: Nat, -I: Array<U32>,
301  -kind: Gather.Lookup, -count: Nat, -inner: Nat, -o: Nat, -i: Nat, -r: Nat> is Data:
302  Summed{-value: F32, -observed: {Model.observe(F32, depth, out, path) == Some{value} : Maybe<&2, F32>},
303    -added: Added(dg, G, di, I, kind, count, inner, o, i, r, count, 0n, 0.0, value)}
304
305# For a balanced gradient array G, a balanced index array I and every split,
306# the slot of the backward build that holds input cell (o, i, r), below the
307# count, holds the gradient words of the cells (o, k, r) whose entry is i,
308# added to +0.0 in increasing k, and +0.0 where no entry is i: PyTorch's
309# one-thread index_add and scatter_add.
310law gather_scatter_cells:
311  for +depth: Nat
312  for +splits: Nat
313  for +kind: Gather.Lookup
314  for +outer: Nat
315  for +rows: Nat
316  for +inner: Nat
317  for +count: Nat
318  for +cells: Nat
319  for +base: Nat
320  for +dg: Nat
321  for -G: Array<F32>
322  for g: Array<F32>
323  for same: {g == G : Array<F32>}
324  for valid: Model.Balanced(F32, dg, G)
325  for +di: Nat
326  for -I: Array<U32>
327  for index: Array<U32>
328  for same_index: {index == I : Array<U32>}
329  for index_valid: Model.Balanced(U32, di, I)
330  for +path: Model.Path(depth)
331  for +o: Nat
332  for +i: Nat
333  for +r: Nat
334  for i_below: {Nat.is_lt(i, rows) == True{} : Bool}
335  for r_below: {Nat.is_lt(r, inner) == True{} : Bool}
336  for at: {Nat.add(Model.offset(depth, path), base) == position(rows, inner, o, i, r) : Nat}
337  for in_count: {Nat.is_lt(Nat.add(Model.offset(depth, path), base), cells) == True{} : Bool}
338  Summed<depth, Gather.scatter_build(depth, splits, Gather.Plan{kind, outer, rows, inner, count}, cells, base, g, index), path,
339    dg, G, di, I, kind, count, inner, o, i, r>