src/tensor/gather.bend

bendtorch-core@0.1.0.0 · 466 lines · 64 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/main.bend as Storage
4import bendtorch-core@0.1.0.0/src/storage/model.bend as Model
5import bendtorch-core@0.1.0.0/src/storage/certificates.bend as Certificates
6import bendtorch-core@0.1.0.0/src/tensor/core.bend as Core
7import bendtorch-core@0.1.0.0/src/tensor/layout.bend as Layout
8import bendtorch-core@0.1.0.0/src/tensor/lifecycle.bend as Lifecycle
9import bendtorch-core@0.1.0.0/src/tensor/parallel.bend as Parallel
10import bendtorch-core@0.1.0.0/src/tensor/slice.bend as Slice
11import bendtorch-core@0.1.0.0/src/tensor/signed.bend as Signed
12
13# PyTorch's index_select(dim, index) and gather(dim, index), as copies (there
14# are no views yet), and their backward, index_add and scatter_add into a
15# +0.0 gradient of the input's shape.
16#
17# There are no integer tensors (every tensor is F32), so an index is a list
18# of U32 entries, never F32 words: index_select takes a vector of entries,
19# and gather an Index, the shape of PyTorch's index tensor and its entries
20# in row-major order. PyTorch refuses a negative entry in both, so
21# a U32 entry loses nothing; an entry at or past the extent is OutOfRange
22# (20). The axis is a signed position, wrapped among the rank: one outside
23# -rank..rank-1, and every axis of a rank-0 input, is AxisOutOfBounds (13).
24#
25# index_select's output is the input with the extent at the axis replaced by
26# the entry count: output cell (o, k, r) of the row-major view [outer, length,
27# inner] holds input cell (o, index[k], r) of [outer, extent, inner].
28# gather's index has the input's rank (else IndexRankMismatch, 21) and the
29# input's extent at every other axis (else IndexShapeMismatch, 22); its
30# entry count must be its shape's element count (else ValueCountMismatch,
31# 10). Its output has the index's dimensions, and output
32# cell (o, k, r) holds input cell (o, index[o, k, r], r).
33#
34# Both builds run on Parallel.build, so no word depends on the split or the
35# thread count. The state of a cell is the source array and the index array.
36# The backward fills input cell (o, i, r) with the gradient words of cells
37# (o, k, r) whose entry is i, added in increasing k to +0.0, the order of
38# PyTorch's one-thread index_add and scatter_add, and +0.0 where no entry is
39# i.
40
41type Error is Data:
42  TensorError{error: Core.Error}
43  LayoutError{error: Layout.Error}
44  OutOfRange{}
45  IndexRankMismatch{}
46  IndexShapeMismatch{}
47
48# PyTorch's index tensor for gather: its shape, a Shape as every tensor's
49# (rank at most 4, at most 2^20 elements), and its entries in row-major
50# order.
51type Index is Data:
52  Index{shape: Shape.Shape, entries: List<&2, U32>}
53
54def built_index(+entries: List<&2, U32>, -dimensions: List<&2, U32>, built: Shape.ConstructResult<dimensions>) -> Result<&2, &2, Error, Index>:
55  match built:
56    case Shape.SError{error}:
57      Fail{TensorError{Core.ShapeError{error}}}
58    case Shape.SBuilt{+shape, preserved}:
59      Done{Index{shape, entries}}
60
61# An Index from an index tensor's dimensions and entries, refused as a tensor
62# of those dimensions is: past rank 4 (RankExceeded, 1) or 2^20 elements
63# (CapacityExceeded, 2). Its entry count is checked by gather.
64def index(+dimensions: List<&2, U32>, +entries: List<&2, U32>) -> Result<&2, &2, Error, Index>:
65  built_index(entries, dimensions, Shape.construct(dimensions))
66
67# Which entry a cell reads: entry k of an index_select, the entry at the
68# output position itself for a gather.
69type Lookup is Data:
70  Selected{}
71  Gathered{}
72
73# What every cell of both builds reads: the kind, the view [outer, extent,
74# inner] of the input around the axis, and the length of the output's axis.
75type Plan is Data:
76  Plan{kind: Lookup, outer: Nat, extent: Nat, inner: Nat, length: Nat}
77
78def plan_length(+plan: Plan) -> Nat:
79  Plan{kind, outer, extent, inner, length} = plan
80  length
81
82def count(entries: List<&2, U32>) -> Nat:
83  match entries:
84    case Nil{}:
85      0n
86    case head <> rest:
87      1n+count(rest)
88
89# Whether every entry is below the extent.
90def below(entries: List<&2, U32>, +extent: Nat) -> Bool:
91  match entries:
92    case Nil{}:
93      True{}
94    case head <> rest:
95      Bool.and(Nat.is_lt(U32.to_nat(head), extent), below(rest, extent))
96
97# Whether two lists of dimensions are equal.
98def equal(left: List<&2, U32>, right: List<&2, U32>) -> Bool:
99  match left right:
100    case Nil{} Nil{}:
101      True{}
102    case x <> xs y <> ys:
103      Bool.and(Nat.is_eq(U32.to_nat(x), U32.to_nat(y)), equal(xs, ys))
104    case _ _:
105      False{}
106
107# Whether the index's dimensions are the input's at every axis but `dim`.
108def matches(dimensions: List<&2, U32>, index: List<&2, U32>, +dim: Nat) -> Bool:
109  match dimensions index dim:
110    case extent <> rest other <> more 0n:
111      equal(rest, more)
112    case extent <> rest other <> more 1n+next:
113      Bool.and(Nat.is_eq(U32.to_nat(extent), U32.to_nat(other)), matches(rest, more, next))
114    case _ _ _:
115      False{}
116
117# The plan of a gather or an index_select of `length` entries at an axis in
118# range.
119def viewed(+kind: Lookup, +given: Slice.View, +length: Nat) -> Plan:
120  Slice.View{+outer, +extent, +inner} = given
121  Plan{kind, outer, extent, inner, length}
122
123def planned_at(+kind: Lookup, +dimensions: List<&2, U32>, +axis: Nat, +length: Nat) -> Plan:
124  viewed(kind, Slice.sized(dimensions, axis, Slice.holds_zero(dimensions)), length)
125
126# A plan the checks accept, and its output shape.
127type Planned is Data:
128  Planned{plan: Plan, output: Shape.Shape}
129
130def shaped(+plan: Plan, -dimensions: List<&2, U32>, built: Shape.ConstructResult<dimensions>) -> Result<&2, &2, Error, Planned>:
131  match built:
132    case Shape.SError{error}:
133      Fail{TensorError{Core.ShapeError{error}}}
134    case Shape.SBuilt{+shape, preserved}:
135      Done{Planned{plan, shape}}
136
137# The output keeps every axis of the input but `axis`, whose extent becomes
138# the length.
139def accepted(+kind: Lookup, +dimensions: List<&2, U32>, +axis: Nat, +length: Nat) -> Result<&2, &2, Error, Planned>:
140  +output = Slice.resized(dimensions, axis, U32.from_nat(length))
141  shaped(planned_at(kind, dimensions, axis, length), output, Shape.construct(output))
142
143def ranged(+kind: Lookup, +dimensions: List<&2, U32>, +axis: Nat, +entries: List<&2, U32>, +length: Nat, inside: Bool) ->
144  Result<&2, &2, Error, Planned>:
145  match inside:
146    case False{}:
147      Fail{OutOfRange{}}
148    case True{}:
149      accepted(kind, dimensions, axis, length)
150
151# The axis check, then the entries against the extent.
152def selecting(+dimensions: List<&2, U32>, +axis: Nat, +entries: List<&2, U32>, in_rank: Bool) -> Result<&2, &2, Error, Planned>:
153  match in_rank:
154    case False{}:
155      Fail{LayoutError{Layout.AxisOutOfBounds{}}}
156    case True{}:
157      ranged(Selected{}, dimensions, axis, entries, count(entries), below(entries, Slice.axis_extent(dimensions, axis)))
158
159# The plan of index_select(dim, index), or its refusal.
160def select_plan(+input: Shape.Shape, dim: Signed.Position, +entries: List<&2, U32>) -> Result<&2, &2, Error, Planned>:
161  +dimensions = Shape.dimensions(input)
162  +axis = U32.to_nat(Signed.wrap(Shape.rank(input), dim))
163  selecting(dimensions, axis, entries, Nat.is_lt(axis, Shape.rank_dimensions(dimensions)))
164
165def gather_shaped(+dimensions: List<&2, U32>, +axis: Nat, +index: List<&2, U32>, +entries: List<&2, U32>, same: Bool) ->
166  Result<&2, &2, Error, Planned>:
167  match same:
168    case False{}:
169      Fail{IndexShapeMismatch{}}
170    case True{}:
171      ranged(Gathered{}, dimensions, axis, entries, Slice.axis_extent(index, axis), below(entries, Slice.axis_extent(dimensions, axis)))
172
173def gather_ranked(+dimensions: List<&2, U32>, +axis: Nat, +index: List<&2, U32>, +entries: List<&2, U32>, same: Bool) ->
174  Result<&2, &2, Error, Planned>:
175  match same:
176    case False{}:
177      Fail{IndexRankMismatch{}}
178    case True{}:
179      gather_shaped(dimensions, axis, index, entries, matches(dimensions, index, axis))
180
181def gather_axis(+dimensions: List<&2, U32>, +axis: Nat, +index: List<&2, U32>, +entries: List<&2, U32>, in_rank: Bool) ->
182  Result<&2, &2, Error, Planned>:
183  match in_rank:
184    case False{}:
185      Fail{LayoutError{Layout.AxisOutOfBounds{}}}
186    case True{}:
187      gather_ranked(dimensions, axis, index, entries, Nat.is_eq(Shape.rank_dimensions(index), Shape.rank_dimensions(dimensions)))
188
189def gather_counted(+dimensions: List<&2, U32>, +axis: Nat, +index: List<&2, U32>, +entries: List<&2, U32>, exact: Bool) ->
190  Result<&2, &2, Error, Planned>:
191  match exact:
192    case False{}:
193      Fail{TensorError{Core.ValueCountMismatch{}}}
194    case True{}:
195      gather_axis(dimensions, axis, index, entries, Nat.is_lt(axis, Shape.rank_dimensions(dimensions)))
196
197# The plan of gather(dim, index), or its refusal, in this order: an entry
198# count other than the index shape's element count, the axis, the index's
199# rank, its extents at the other axes, then its entries. The element count
200# is the Shape's own, so no product of its extents is formed here.
201def gather_plan(+input: Shape.Shape, dim: Signed.Position, +index: Index) -> Result<&2, &2, Error, Planned>:
202  Index{+shape, +entries} = index
203  +axis = U32.to_nat(Signed.wrap(Shape.rank(input), dim))
204  gather_counted(Shape.dimensions(input), axis, Shape.dimensions(shape), entries,
205    Nat.is_eq(count(entries), Shape.elements(shape)))
206
207# The entries as a balanced U32 array that a cell reads in a few steps: the
208# smallest depth that holds them, as Storage allocates.
209def written(entries: List<&2, U32>, array: Array<U32>, +at: U32) -> Array<U32>:
210  match entries:
211    case Nil{}:
212      array
213    case head <> rest:
214      written(rest, Array.set(U32, array, at, head), U32.add(at, 1))
215
216def index_depth(pair: Nat & Nat) -> Nat:
217  (depth, capacity) = pair
218  depth
219
220def index_array(+entries: List<&2, U32>) -> Array<U32>:
221  +n = count(entries)
222  written(entries, Array.new(U32, index_depth(Storage.depth_seek(20n, n, 0n, 1n, Nat.is_le(n, 1n))), 0), 0)
223
224# The cells of the forward build.
225
226# The position of cell (o, k, r) of the view [outer, rows, inner].
227def position(+rows: Nat, +inner: Nat, +o: Nat, +k: Nat, +r: Nat) -> Nat:
228  Nat.add(Nat.mul(Nat.add(Nat.mul(o, rows), k), inner), r)
229
230# The entry output cell (o, k, r) reads: k for an index_select, the cell's
231# own position for a gather.
232def entry_of(+kind: Lookup, +length: Nat, +inner: Nat, +o: Nat, +k: Nat, +r: Nat) -> Nat:
233  match kind:
234    case Selected{}:
235      k
236    case Gathered{}:
237      position(length, inner, o, k, r)
238
239# Output cell `at` is (o, k, r) of [outer, length, inner]; with entry j, it
240# reads input cell (o, j, r) of [outer, extent, inner].
241def entry(+plan: Plan, +at: Nat) -> Nat:
242  Plan{+kind, +outer, +extent, +inner, +length} = plan
243  +row = Nat.div(at, inner)
244  entry_of(kind, length, inner, Nat.div(row, length), Nat.mod(row, length), Nat.mod(at, inner))
245
246def source(+plan: Plan, +at: Nat, +j: Nat) -> Nat:
247  Plan{+kind, +outer, +extent, +inner, +length} = plan
248  position(extent, inner, Nat.div(Nat.div(at, inner), length), j, Nat.mod(at, inner))
249
250def same(+plan: Plan) -> Plan:
251  plan
252
253def take(index: Array<U32>, pair: Array<F32> & F32) -> F32 & (Array<F32> & Array<U32>):
254  (a, x) = pair
255  (x, (a, index))
256
257def looked(+plan: Plan, +at: Nat, a: Array<F32>, pair: Array<U32> & U32) -> F32 & (Array<F32> & Array<U32>):
258  (index, +j) = pair
259  take(index, Array.get(F32, a, U32.from_nat(source(plan, at, U32.to_nat(j)))))
260
261def read(+plan: Plan, +at: Nat, state: Array<F32> & Array<U32>) -> F32 & (Array<F32> & Array<U32>):
262  (a, index) = state
263  looked(plan, at, a, Array.get(U32, index, U32.from_nat(entry(plan, at))))
264
265# A cell past the count is +0.0 padding.
266def cell(+plan: Plan, +at: Nat, state: Array<F32> & Array<U32>, inside: Bool) -> F32 & (Array<F32> & Array<U32>):
267  match inside:
268    case False{}:
269      (0.0, state)
270    case True{}:
271      read(plan, at, state)
272
273def halves(a: Array<F32> & Array<F32>, b: Array<U32> & Array<U32>) ->
274  (Array<F32> & Array<U32>) & (Array<F32> & Array<U32>):
275  (a1, a2) = a
276  (b1, b2) = b
277  ((a1, b1), (a2, b2))
278
279# Each half of a fork gets its own copy of both arrays.
280def split(state: Array<F32> & Array<U32>) -> (Array<F32> & Array<U32>) & (Array<F32> & Array<U32>):
281  (a, b) = state
282  halves(Array.clone(F32, a), Array.clone(U32, b))
283
284def build(+depth: Nat, +splits: Nat, +plan: Plan, +count: Nat, +base: Nat, a: Array<F32>, index: Array<U32>) -> Array<F32>:
285  Parallel.build(~Plan, ~Plan, ~Array<F32>, ~Array<U32>, ~same, ~split, ~cell, plan, depth, splits, count, base, (a, index))
286
287def balanced(+depth: Nat, +splits: Nat, +plan: Plan, +count: Nat, +base: Nat, a: Array<F32>, index: Array<U32>) ->
288  Model.Balanced(F32, depth, build(depth, splits, plan, count, base, a, index)):
289  Parallel.balanced(~Plan, ~Plan, ~Array<F32>, ~Array<U32>, ~same, ~split, ~cell, plan, depth, splits, count, base, (a, index))
290
291# The cells of the backward build, over the input's shape.
292
293def accumulate(acc: F32, index: Array<U32>, pair: Array<F32> & F32) -> F32 & (Array<F32> & Array<U32>):
294  (g, x) = pair
295  (F32.add(acc, x), (g, index))
296
297def summed(hit: Bool, +at: Nat, acc: F32, g: Array<F32>, index: Array<U32>) -> F32 & (Array<F32> & Array<U32>):
298  match hit:
299    case False{}:
300      (acc, (g, index))
301    case True{}:
302      accumulate(acc, index, Array.get(F32, g, U32.from_nat(at)))
303
304def compared(+i: Nat, +at: Nat, acc: F32, g: Array<F32>, pair: Array<U32> & U32) -> F32 & (Array<F32> & Array<U32>):
305  (index, +j) = pair
306  summed(Nat.is_eq(U32.to_nat(j), i), at, acc, g, index)
307
308# One gradient cell (o, k, r): its word is added when its entry is i.
309def stepped(+i: Nat, +at: Nat, +slot: Nat, acc: F32, state: Array<F32> & Array<U32>) -> F32 & (Array<F32> & Array<U32>):
310  (g, index) = state
311  compared(i, at, acc, g, Array.get(U32, index, U32.from_nat(slot)))
312
313# The gradient cells (o, k, r), (o, k + 1, r), ... for `left` more k, added
314# to acc in increasing k.
315# Gradient cell (o, k, r) and the index slot it reads.
316def cell_at(+plan: Plan, +o: Nat, +k: Nat, +r: Nat) -> Nat:
317  Plan{+kind, +outer, +extent, +inner, +length} = plan
318  position(length, inner, o, k, r)
319
320def slot_at(+plan: Plan, +o: Nat, +k: Nat, +r: Nat) -> Nat:
321  Plan{+kind, +outer, +extent, +inner, +length} = plan
322  entry_of(kind, length, inner, o, k, r)
323
324def stepped_pair(+i: Nat, +at: Nat, +slot: Nat, pair: F32 & (Array<F32> & Array<U32>)) -> F32 & (Array<F32> & Array<U32>):
325  (acc, state) = pair
326  stepped(i, at, slot, acc, state)
327
328def added(+plan: Plan, +o: Nat, +i: Nat, +r: Nat, left: Nat, +k: Nat, pair: F32 & (Array<F32> & Array<U32>)) ->
329  F32 & (Array<F32> & Array<U32>):
330  match left:
331    case 0n:
332      pair
333    case 1n+rest:
334      added(plan, o, i, r, rest, 1n+k, stepped_pair(i, cell_at(plan, o, k, r), slot_at(plan, o, k, r), pair))
335
336# Input cell `at` is (o, i, r) of [outer, extent, inner].
337def scattered(+plan: Plan, +at: Nat, state: Array<F32> & Array<U32>) -> F32 & (Array<F32> & Array<U32>):
338  Plan{+kind, +outer, +extent, +inner, +length} = plan
339  +row = Nat.div(at, inner)
340  added(plan, Nat.div(row, extent), Nat.mod(row, extent), Nat.mod(at, inner), length, 0n, (0.0, state))
341
342def scatter_cell(+plan: Plan, +at: Nat, state: Array<F32> & Array<U32>, inside: Bool) -> F32 & (Array<F32> & Array<U32>):
343  match inside:
344    case False{}:
345      (0.0, state)
346    case True{}:
347      scattered(plan, at, state)
348
349def scatter_build(+depth: Nat, +splits: Nat, +plan: Plan, +count: Nat, +base: Nat, g: Array<F32>, index: Array<U32>) -> Array<F32>:
350  Parallel.build(~Plan, ~Plan, ~Array<F32>, ~Array<U32>, ~same, ~split, ~scatter_cell, plan, depth, splits, count, base, (g, index))
351
352def scatter_balanced(+depth: Nat, +splits: Nat, +plan: Plan, +count: Nat, +base: Nat, g: Array<F32>, index: Array<U32>) ->
353  Model.Balanced(F32, depth, scatter_build(depth, splits, plan, count, base, g, index)):
354  Parallel.balanced(~Plan, ~Plan, ~Array<F32>, ~Array<U32>, ~same, ~split, ~scatter_cell, plan, depth, splits, count, base, (g, index))
355
356# Running a build on a tensor.
357
358type Pass is Data:
359  Forward{}
360  Backward{}
361
362def built_array(kind: Pass, +depth: Nat, +plan: Plan, +count: Nat, a: Array<F32>, index: Array<U32>) -> Array<F32>:
363  match kind:
364    case Forward{}:
365      build(depth, Slice.splits(depth), plan, count, 0n, a, index)
366    case Backward{}:
367      scatter_build(depth, Slice.splits(depth), plan, count, 0n, a, index)
368
369def built_balanced(kind: Pass, +depth: Nat, +plan: Plan, +count: Nat, a: Array<F32>, index: Array<U32>) ->
370  Model.Balanced(F32, depth, built_array(kind, depth, plan, count, a, index)):
371  match kind:
372    case Forward{}:
373      balanced(depth, Slice.splits(depth), plan, count, 0n, a, index)
374    case Backward{}:
375      scatter_balanced(depth, Slice.splits(depth), plan, count, 0n, a, index)
376
377def built_cloned(+kind: Pass, +input: Shape.Shape, +output: Shape.Shape, +plan: Plan, +entries: List<&2, U32>, +depth: Nat, +length: Nat,
378  -valid: {Storage.bounds(depth, length) == True{} : Bool},
379  -exact: {length == Shape.logical_numel(Shape.dimensions(output)) : Nat},
380  +d: Nat, +n: Nat, -array: Array<F32>, pair: Array<F32> & Array<F32>, -same: {pair == Array.clone(F32, array) : Array<F32> & Array<F32>},
381  -balanced: Model.Balanced(F32, d, array), -ok: {Storage.bounds(d, n) == True{} : Bool},
382  -fits: {Lifecycle.logical_length(Storage.Owned{d, n, array, balanced, ok}) == Shape.logical_numel(Shape.dimensions(input)) : Nat}) ->
383  Core.Dynamic & Result<&2, &1, Error, Core.Dynamic>:
384  (kept, copy) = pair
385  (Core.Packed{input, Core.TTensor{Storage.Owned{d, n, kept, Certificates.clone_first(F32, d, array, (kept, copy), same, balanced), ok}, fits}},
386    Done{Core.Packed{output, Core.TTensor{Storage.Owned{depth, length, built_array(kind, depth, plan, length, copy, index_array(entries)),
387      built_balanced(kind, depth, plan, length, copy, index_array(entries)), valid}, exact}}})
388
389# An empty source: the forward copy never has cells then, and the backward
390# of an empty gradient reads nothing.
391def built_empty(+kind: Pass, +output: Shape.Shape, +plan: Plan, +entries: List<&2, U32>, +depth: Nat, +length: Nat,
392  -valid: {Storage.bounds(depth, length) == True{} : Bool},
393  -exact: {length == Shape.logical_numel(Shape.dimensions(output)) : Nat}) -> Result<&2, &1, Error, Core.Dynamic>:
394  match kind:
395    case Forward{}:
396      Fail{TensorError{Core.StorageError{Storage.OutOfBounds{}}}}
397    case Backward{}:
398      Done{Core.Packed{output, Core.TTensor{Storage.Owned{depth, length,
399        built_array(Backward{}, depth, plan, length, Array.new(F32, 0n, 0.0), index_array(entries)),
400        built_balanced(Backward{}, depth, plan, length, Array.new(F32, 0n, 0.0), index_array(entries)), valid}, exact}}}
401
402def built_owned(+kind: Pass, +input: Shape.Shape, +output: Shape.Shape, +plan: Plan, +entries: List<&2, U32>, +depth: Nat, +length: Nat,
403  -valid: {Storage.bounds(depth, length) == True{} : Bool},
404  -exact: {length == Shape.logical_numel(Shape.dimensions(output)) : Nat},
405  source: Core.Tensor<input>) -> Core.Dynamic & Result<&2, &1, Error, Core.Dynamic>:
406  Core.TTensor{storage, fits} = source
407  match storage:
408    case Storage.Empty{}:
409      (Core.Packed{input, Core.TTensor{Storage.Empty{}, fits}}, built_empty(kind, output, plan, entries, depth, length, valid, exact))
410    case Storage.Owned{+d, +n, array, balanced, ok}:
411      built_cloned(kind, input, output, plan, entries, depth, length, valid, exact, d, n, array, Array.clone(F32, array), {==}, balanced, ok, fits)
412
413# Every output cell at once, into the storage that `Core.new` allocates. An
414# output with no cell keeps that storage, which is empty.
415def built(+kind: Pass, +input: Shape.Shape, +output: Shape.Shape, +plan: Plan, +entries: List<&2, U32>, source: Core.Tensor<input>,
416  allocated: Result<&2, &1, Core.Error, Core.Tensor<output>>) -> Core.Dynamic & Result<&2, &1, Error, Core.Dynamic>:
417  match allocated:
418    case Fail{error}:
419      (Core.Packed{input, source}, Fail{TensorError{error}})
420    case Done{target}:
421      Core.TTensor{storage, exact} = target
422      match storage:
423        case Storage.Empty{}:
424          (Core.Packed{input, source}, Done{Core.Packed{output, Core.TTensor{Storage.Empty{}, exact}}})
425        case Storage.Owned{+depth, +length, array, balanced, valid}:
426          built_owned(kind, input, output, plan, entries, depth, length, valid, exact, source)
427
428def copied(+shape: Shape.Shape, tensor: Core.Tensor<shape>, +entries: List<&2, U32>, planned: Result<&2, &2, Error, Planned>) ->
429  Core.Dynamic & Result<&2, &1, Error, Core.Dynamic>:
430  match planned:
431    case Fail{error}:
432      (Core.Packed{shape, tensor}, Fail{error})
433    case Done{Planned{+plan, +output}}:
434      built(Forward{}, shape, output, plan, entries, tensor, Core.new(output, 0.0))
435
436# PyTorch's input.index_select(dim, index). The input owner always comes back
437# unchanged next to the result.
438def index_select(input: Core.Dynamic, dim: Signed.Position, +entries: List<&2, U32>) -> Core.Dynamic & Result<&2, &1, Error, Core.Dynamic>:
439  Core.Packed{+shape, tensor} = input
440  copied(shape, tensor, entries, select_plan(shape, dim, entries))
441
442# PyTorch's input.gather(dim, index).
443def gather(input: Core.Dynamic, dim: Signed.Position, +index: Index) -> Core.Dynamic & Result<&2, &1, Error, Core.Dynamic>:
444  Core.Packed{+shape, tensor} = input
445  Index{+index_shape, +entries} = index
446  copied(shape, tensor, entries, gather_plan(shape, dim, Index{index_shape, entries}))
447
448# The backward of a planned index_select or gather: the gradient g of its
449# output, added into +0.0 words of the input's shape at the entries, PyTorch's
450# index_add and scatter_add. g comes back unchanged next to the result.
451def scatter(g: Core.Dynamic, +input: Shape.Shape, +plan: Plan, +entries: List<&2, U32>) -> Core.Dynamic & Result<&2, &1, Error, Core.Dynamic>:
452  Core.Packed{+shape, tensor} = g
453  built(Backward{}, shape, input, plan, entries, tensor, Core.new(input, 0.0))
454
455def error_code(error: Error) -> U32:
456  match error:
457    case TensorError{error}:
458      Core.error_code(error)
459    case LayoutError{error}:
460      Slice.error_code(Slice.LayoutError{error})
461    case OutOfRange{}:
462      20
463    case IndexRankMismatch{}:
464      21
465    case IndexShapeMismatch{}:
466      22