1	import Base
2	import ../shape/core.bend as Shape
3	import ../storage/main.bend as Storage
4	import ../storage/model.bend as Model
5	import ../storage/certificates.bend as Certificates
6	import ./core.bend as Core
7	import ./layout.bend as Layout
8	import ./lifecycle.bend as Lifecycle
9	import ./parallel.bend as Parallel
10	import ./slice.bend as Slice
11	import ./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	
41	type 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.
51	type Index is Data:
52	  Index{shape: Shape.Shape, entries: List<&2, U32>}
53	
54	def 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.
64	def 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.
69	type 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.
75	type Plan is Data:
76	  Plan{kind: Lookup, outer: Nat, extent: Nat, inner: Nat, length: Nat}
77	
78	def plan_length(+plan: Plan) -> Nat:
79	  Plan{kind, outer, extent, inner, length} = plan
80	  length
81	
82	def 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.
90	def 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.
98	def 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`.
108	def 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.
119	def viewed(+kind: Lookup, +given: Slice.View, +length: Nat) -> Plan:
120	  Slice.View{+outer, +extent, +inner} = given
121	  Plan{kind, outer, extent, inner, length}
122	
123	def 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.
127	type Planned is Data:
128	  Planned{plan: Plan, output: Shape.Shape}
129	
130	def 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.
139	def 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	
143	def 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.
152	def 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.
160	def 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	
165	def 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	
173	def 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	
181	def 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	
189	def 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.
201	def 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.
209	def 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	
216	def index_depth(pair: Nat & Nat) -> Nat:
217	  (depth, capacity) = pair
218	  depth
219	
220	def 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].
227	def 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.
232	def 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].
241	def 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	
246	def 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	
250	def same(+plan: Plan) -> Plan:
251	  plan
252	
253	def take(index: Array<U32>, pair: Array<F32> & F32) -> F32 & (Array<F32> & Array<U32>):
254	  (a, x) = pair
255	  (x, (a, index))
256	
257	def 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	
261	def 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.
266	def 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	
273	def 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.
280	def 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	
284	def 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	
287	def 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	
293	def 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	
297	def 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	
304	def 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.
309	def 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.
316	def 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	
320	def 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	
324	def 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	
328	def 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].
337	def 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	
342	def 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	
349	def 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	
352	def 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	
358	type Pass is Data:
359	  Forward{}
360	  Backward{}
361	
362	def 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	
369	def 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	
377	def 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.
391	def 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	
402	def 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.
415	def 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	
428	def 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.
438	def 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).
443	def 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.
451	def 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	
455	def 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