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 ./transforms.bend as Transforms
10	import ./parallel.bend as Parallel
11	import ./elementwise.bend as Elementwise
12	
13	# PyTorch's basic slicing along one axis, as a copy (there are no views
14	# yet): `input[..., start:stop:step, ...]` with the slice at axis `dim`
15	# (aten::slice), `narrow(dim, start, length)` and `select(dim, index)`, which
16	# drops the axis. A multi-axis slice `t[a:b, c:d]` is one slice per axis, as
17	# PyTorch's indexing applies them.
18	#
19	# The axis is checked first: a rank-0 input has no axis to slice, and an axis
20	# at or past the rank is AxisOutOfBounds (13), permute's code. Then a step of
21	# 0 is StepNotPositive (19), as PyTorch refuses every step <= 0. The start and
22	# stop are clamped to the axis as PyTorch clamps them (`bounds`), so a stop
23	# past the extent ends at the extent and a stop before the start gives an
24	# empty axis. `narrow` refuses a start past the extent or a start + length
25	# past it, and `select` an index at or past the extent, with OutOfRange (20).
26	#
27	# Every index here is a U32, so none is negative. PyTorch's signed axes and
28	# positions (-1 is the last axis or element) are parsed in signed.bend
29	# (Signed.slice, narrow and select), which calls these operations; the tape
30	# records slices through those. `bounds` is the one place a start and a stop
31	# are clamped to the extent.
32	#
33	# Output cell (o, k, r) of the row-major view [outer, length, inner] around
34	# the axis holds input cell (o, start + k * step, r) of [outer, extent,
35	# inner]. The cells are built by `Parallel.build`, so no word depends on the
36	# split or the thread count. The backward, PyTorch's slice_backward, scatters
37	# a gradient of the output's shape into +0.0 words of the input's shape: input
38	# cell (o, i, r) holds gradient cell (o, k, r) where i = start + k * step for
39	# some k below the length, and +0.0 everywhere else.
40	
41	type Error is Data:
42	  TensorError{error: Core.Error}
43	  LayoutError{error: Layout.Error}
44	  StepNotPositive{}
45	  OutOfRange{}
46	
47	# A slice's start and stop after PyTorch's clamp to an axis of extent n:
48	# start' = min(start, n), stop' = max(start', min(stop, n)).
49	type Bounds is Data:
50	  Bounds{start: Nat, stop: Nat}
51	
52	def bounds(+extent: Nat, +start: Nat, +stop: Nat) -> Bounds:
53	  +low = Nat.min(start, extent)
54	  Bounds{low, Nat.max(low, Nat.min(stop, extent))}
55	
56	# How many of 0, step, 2 step, ... lie below `span`: none for an empty span,
57	# else 1 + (span - 1) / step, a ceiling division that never forms a number
58	# past span. It takes one division, not a step per index.
59	def counted(+span: Nat, +step: Nat) -> Nat:
60	  match span:
61	    case 0n:
62	      0n
63	    case 1n+rest:
64	      1n+Nat.div(rest, step)
65	
66	# The length of the sliced axis: the number of indices start', start' +
67	# step, ... below stop'.
68	def length(+clamped: Bounds, +step: Nat) -> Nat:
69	  Bounds{+low, +high} = clamped
70	  counted(Nat.sub(high, low), step)
71	
72	# The row-major view of the input around the axis: the element count before
73	# it, its extent and the element count after it.
74	type View is Data:
75	  View{outer: Nat, extent: Nat, inner: Nat}
76	
77	def view_at(dimensions: List<&2, U32>, +dim: Nat, +outer: Nat) -> View:
78	  match dimensions dim:
79	    case extent <> rest 0n:
80	      View{outer, U32.to_nat(extent), Shape.logical_numel(rest)}
81	    case extent <> rest 1n+next:
82	      view_at(rest, next, Nat.mul(outer, U32.to_nat(extent)))
83	    case Nil{} _:
84	      View{outer, 0n, 1n}
85	
86	def view(dimensions: List<&2, U32>, +dim: Nat) -> View:
87	  view_at(dimensions, dim, 1n)
88	
89	# The extent of axis `dim`, 0 past the rank. It reads that one dimension and
90	# forms no product.
91	def axis_extent(dimensions: List<&2, U32>, +dim: Nat) -> Nat:
92	  match dimensions dim:
93	    case extent <> rest 0n:
94	      U32.to_nat(extent)
95	    case extent <> rest 1n+next:
96	      axis_extent(rest, next)
97	    case Nil{} _:
98	      0n
99	
100	def extent_at(dimensions: List<&2, U32>, +dim: U32) -> Nat:
101	  axis_extent(dimensions, U32.to_nat(dim))
102	
103	# Whether some dimension is 0, so the input has no element.
104	def holds_zero(dimensions: List<&2, U32>) -> Bool:
105	  match dimensions:
106	    case Nil{}:
107	      False{}
108	    case extent <> rest:
109	      Bool.or(Nat.is_eq(U32.to_nat(extent), 0n), holds_zero(rest))
110	
111	# The view a plan uses. An input with no element gets no cell in either
112	# build, so its outer and inner counts are left at 0 rather than formed: the
113	# product of the other extents can pass 2^48 - 1, the largest Nat the
114	# runtime holds, as [2^24, 2^24, 0] at axis 2 does.
115	def sized(+dimensions: List<&2, U32>, +dim: Nat, empty: Bool) -> View:
116	  match empty:
117	    case True{}:
118	      View{0n, axis_extent(dimensions, dim), 0n}
119	    case False{}:
120	      view(dimensions, dim)
121	
122	# The input dimensions with the extent at the axis replaced.
123	def resized(dimensions: List<&2, U32>, +dim: Nat, +extent: U32) -> List<&2, U32>:
124	  match dimensions dim:
125	    case head <> rest 0n:
126	      extent <> rest
127	    case head <> rest 1n+next:
128	      head <> resized(rest, next, extent)
129	    case Nil{} _:
130	      Nil{}
131	
132	# The input dimensions without the axis, as `select` returns them.
133	def dropped(dimensions: List<&2, U32>, +dim: Nat) -> List<&2, U32>:
134	  match dimensions dim:
135	    case head <> rest 0n:
136	      rest
137	    case head <> rest 1n+next:
138	      head <> dropped(rest, next)
139	    case Nil{} _:
140	      Nil{}
141	
142	# What every cell of a slice and of its backward reads: the view, the
143	# clamped start, the step and the length of the sliced axis.
144	type Plan is Data:
145	  Plan{outer: Nat, extent: Nat, inner: Nat, start: Nat, step: Nat, length: Nat}
146	
147	def planned_view(+outer: Nat, +extent: Nat, +inner: Nat, +clamped: Bounds, +step: Nat) -> Plan:
148	  Bounds{+low, +high} = clamped
149	  Plan{outer, extent, inner, low, step, length(clamped, step)}
150	
151	def viewed(+given: View, +start: Nat, +stop: Nat, +step: Nat) -> Plan:
152	  View{+outer, +extent, +inner} = given
153	  planned_view(outer, extent, inner, bounds(extent, start, stop), step)
154	
155	def plan_of(+dimensions: List<&2, U32>, +dim: U32, +start: U32, +stop: U32, +step: U32) -> Plan:
156	  viewed(sized(dimensions, U32.to_nat(dim), holds_zero(dimensions)), U32.to_nat(start), U32.to_nat(stop), U32.to_nat(step))
157	
158	def plan_length(+plan: Plan) -> Nat:
159	  Plan{outer, extent, inner, start, step, length} = plan
160	  length
161	
162	# PyTorch's checks, in its order: the axis (a rank-0 input has none), then
163	# the step.
164	def refused(inside: Bool, positive: Bool) -> Maybe<&2, Error>:
165	  match inside positive:
166	    case False{} _:
167	      Some{LayoutError{Layout.AxisOutOfBounds{}}}
168	    case True{} False{}:
169	      Some{StepNotPositive{}}
170	    case True{} True{}:
171	      None{}
172	
173	def refusal(dimensions: List<&2, U32>, +dim: U32, +step: U32) -> Maybe<&2, Error>:
174	  refused(Nat.is_lt(U32.to_nat(dim), Shape.rank_dimensions(dimensions)), Nat.is_lt(0n, U32.to_nat(step)))
175	
176	# A slice the checks accept: its plan and its output shape.
177	type Planned is Data:
178	  Planned{plan: Plan, output: Shape.Shape}
179	
180	def shaped(+plan: Plan, -dimensions: List<&2, U32>, built: Shape.ConstructResult<dimensions>) -> Result<&2, &2, Error, Planned>:
181	  match built:
182	    case Shape.SError{error}:
183	      Fail{TensorError{Core.ShapeError{error}}}
184	    case Shape.SBuilt{+shape, preserved}:
185	      Done{Planned{plan, shape}}
186	
187	def checked(+dimensions: List<&2, U32>, +dim: U32, +start: U32, +stop: U32, +step: U32, refused: Maybe<&2, Error>) ->
188	  Result<&2, &2, Error, Planned>:
189	  match refused:
190	    case Some{error}:
191	      Fail{error}
192	    case None{}:
193	      +plan = plan_of(dimensions, dim, start, stop, step)
194	      +output = resized(dimensions, U32.to_nat(dim), U32.from_nat(plan_length(plan)))
195	      shaped(plan, output, Shape.construct(output))
196	
197	# The plan of input[..., start:stop:step, ...] at axis `dim`, or its refusal.
198	def plan(+input: Shape.Shape, +dim: U32, +start: U32, +stop: U32, +step: U32) -> Result<&2, &2, Error, Planned>:
199	  checked(Shape.dimensions(input), dim, start, stop, step, refusal(Shape.dimensions(input), dim, step))
200	
201	# The cells of the forward build.
202	
203	# The input index output cell `position` reads: position is (o * length + k)
204	# * inner + r, and it reads (o * extent + start + k * step) * inner + r.
205	def source(+plan: Plan, +position: Nat) -> Nat:
206	  Plan{+outer, +extent, +inner, +start, +step, +length} = plan
207	  +row = Nat.div(position, inner)
208	  +o = Nat.div(row, length)
209	  Nat.add(Nat.mul(Nat.add(Nat.mul(o, extent), Nat.add(start, Nat.mul(Nat.mod(row, length), step))), inner), Nat.mod(position, inner))
210	
211	def same(+plan: Plan) -> Plan:
212	  plan
213	
214	def take(b: Array<F32>, pair: Array<F32> & F32) -> F32 & (Array<F32> & Array<F32>):
215	  (a, x) = pair
216	  (x, (a, b))
217	
218	def read(+plan: Plan, +position: Nat, state: Array<F32> & Array<F32>) -> F32 & (Array<F32> & Array<F32>):
219	  (a, b) = state
220	  take(b, Array.get(F32, a, U32.from_nat(source(plan, position))))
221	
222	# The state is the input array and an empty array that fills the second
223	# slot of the state pair; a cell past the count is +0.0 padding.
224	def cell(+plan: Plan, +position: Nat, state: Array<F32> & Array<F32>, inside: Bool) -> F32 & (Array<F32> & Array<F32>):
225	  match inside:
226	    case False{}:
227	      (0.0, state)
228	    case True{}:
229	      read(plan, position, state)
230	
231	# A subtree below the forks holds 2^15 cells, as in the other builds.
232	def splits(+depth: Nat) -> Nat:
233	  Nat.sub(depth, 15n)
234	
235	def build(+depth: Nat, +splits: Nat, +plan: Plan, +count: Nat, +base: Nat, a: Array<F32>) -> Array<F32>:
236	  Parallel.build(~Plan, ~Plan, ~Array<F32>, ~Array<F32>, ~same, ~Elementwise.split, ~cell, plan, depth, splits, count, base,
237	    (a, Array.new(F32, 0n, 0.0)))
238	
239	def balanced(+depth: Nat, +splits: Nat, +plan: Plan, +count: Nat, +base: Nat, a: Array<F32>) ->
240	  Model.Balanced(F32, depth, build(depth, splits, plan, count, base, a)):
241	  Parallel.balanced(~Plan, ~Plan, ~Array<F32>, ~Array<F32>, ~same, ~Elementwise.split, ~cell, plan, depth, splits, count, base,
242	    (a, Array.new(F32, 0n, 0.0)))
243	
244	# The cells of the backward build, over the input's shape.
245	
246	# The gradient index input cell `position` reads, or None where no output
247	# cell came from it: position is (o * extent + i) * inner + r, and it reads
248	# (o * length + k) * inner + r when i = start + k * step with k below the
249	# length, for k = (i - start) / step.
250	def picked(+plan: Plan, +row: Nat, +within: Nat, +i: Nat, +k: Nat, hit: Bool) -> Maybe<&2, Nat>:
251	  Plan{+outer, +extent, +inner, +start, +step, +length} = plan
252	  match hit:
253	    case False{}:
254	      None{}
255	    case True{}:
256	      Some{Nat.add(Nat.mul(Nat.add(Nat.mul(Nat.div(row, extent), length), k), inner), within)}
257	
258	def hits(+start: Nat, +step: Nat, +length: Nat, +i: Nat, +k: Nat) -> Bool:
259	  Bool.and(Nat.is_le(start, i), Bool.and(Nat.is_lt(k, length), Nat.is_eq(Nat.add(start, Nat.mul(k, step)), i)))
260	
261	def target(+plan: Plan, +position: Nat) -> Maybe<&2, Nat>:
262	  Plan{+outer, +extent, +inner, +start, +step, +length} = plan
263	  +row = Nat.div(position, inner)
264	  +i = Nat.mod(row, extent)
265	  +k = Nat.div(Nat.sub(i, start), step)
266	  picked(plan, row, Nat.mod(position, inner), i, k, hits(start, step, length, i, k))
267	
268	def fetched(+index: Nat, state: Array<F32> & Array<F32>) -> F32 & (Array<F32> & Array<F32>):
269	  (a, b) = state
270	  take(b, Array.get(F32, a, U32.from_nat(index)))
271	
272	def scattered(state: Array<F32> & Array<F32>, at: Maybe<&2, Nat>) -> F32 & (Array<F32> & Array<F32>):
273	  match at:
274	    case None{}:
275	      (0.0, state)
276	    case Some{+index}:
277	      fetched(index, state)
278	
279	def scatter_cell(+plan: Plan, +position: Nat, state: Array<F32> & Array<F32>, inside: Bool) -> F32 & (Array<F32> & Array<F32>):
280	  match inside:
281	    case False{}:
282	      (0.0, state)
283	    case True{}:
284	      scattered(state, target(plan, position))
285	
286	def scatter_build(+depth: Nat, +splits: Nat, +plan: Plan, +count: Nat, +base: Nat, g: Array<F32>) -> Array<F32>:
287	  Parallel.build(~Plan, ~Plan, ~Array<F32>, ~Array<F32>, ~same, ~Elementwise.split, ~scatter_cell, plan, depth, splits, count, base,
288	    (g, Array.new(F32, 0n, 0.0)))
289	
290	def scatter_balanced(+depth: Nat, +splits: Nat, +plan: Plan, +count: Nat, +base: Nat, g: Array<F32>) ->
291	  Model.Balanced(F32, depth, scatter_build(depth, splits, plan, count, base, g)):
292	  Parallel.balanced(~Plan, ~Plan, ~Array<F32>, ~Array<F32>, ~same, ~Elementwise.split, ~scatter_cell, plan, depth, splits, count, base,
293	    (g, Array.new(F32, 0n, 0.0)))
294	
295	# Running a build on a tensor.
296	
297	# Which build fills the output: the forward copy or the backward scatter.
298	type Pass is Data:
299	  Forward{}
300	  Backward{}
301	
302	def built_array(kind: Pass, +depth: Nat, +plan: Plan, +count: Nat, a: Array<F32>) -> Array<F32>:
303	  match kind:
304	    case Forward{}:
305	      build(depth, splits(depth), plan, count, 0n, a)
306	    case Backward{}:
307	      scatter_build(depth, splits(depth), plan, count, 0n, a)
308	
309	def built_balanced(kind: Pass, +depth: Nat, +plan: Plan, +count: Nat, a: Array<F32>) ->
310	  Model.Balanced(F32, depth, built_array(kind, depth, plan, count, a)):
311	  match kind:
312	    case Forward{}:
313	      balanced(depth, splits(depth), plan, count, 0n, a)
314	    case Backward{}:
315	      scatter_balanced(depth, splits(depth), plan, count, 0n, a)
316	
317	# The build from an owned source: the source comes back as the first half of
318	# its clone, and the build reads the second half.
319	def built_cloned(+kind: Pass, +input: Shape.Shape, +output: Shape.Shape, +plan: Plan, +depth: Nat, +length: Nat,
320	  -valid: {Storage.bounds(depth, length) == True{} : Bool},
321	  -exact: {length == Shape.logical_numel(Shape.dimensions(output)) : Nat},
322	  +d: Nat, +n: Nat, -array: Array<F32>, pair: Array<F32> & Array<F32>, -same: {pair == Array.clone(F32, array) : Array<F32> & Array<F32>},
323	  -balanced: Model.Balanced(F32, d, array), -ok: {Storage.bounds(d, n) == True{} : Bool},
324	  -fits: {Lifecycle.logical_length(Storage.Owned{d, n, array, balanced, ok}) == Shape.logical_numel(Shape.dimensions(input)) : Nat}) ->
325	  Core.Dynamic & Result<&2, &1, Error, Core.Dynamic>:
326	  (kept, copy) = pair
327	  (Core.Packed{input, Core.TTensor{Storage.Owned{d, n, kept, Certificates.clone_first(F32, d, array, (kept, copy), same, balanced), ok}, fits}},
328	    Done{Core.Packed{output, Core.TTensor{Storage.Owned{depth, length, built_array(kind, depth, plan, length, copy),
329	      built_balanced(kind, depth, plan, length, copy), valid}, exact}}})
330	
331	# An empty source: the forward copy never has cells then, and the backward
332	# of an empty gradient reads nothing, so the build reads an empty array.
333	def built_empty(+kind: Pass, +output: Shape.Shape, +plan: Plan, +depth: Nat, +length: Nat,
334	  -valid: {Storage.bounds(depth, length) == True{} : Bool},
335	  -exact: {length == Shape.logical_numel(Shape.dimensions(output)) : Nat}) -> Result<&2, &1, Error, Core.Dynamic>:
336	  match kind:
337	    case Forward{}:
338	      Fail{TensorError{Core.StorageError{Storage.OutOfBounds{}}}}
339	    case Backward{}:
340	      Done{Core.Packed{output, Core.TTensor{Storage.Owned{depth, length, built_array(Backward{}, depth, plan, length, Array.new(F32, 0n, 0.0)),
341	        built_balanced(Backward{}, depth, plan, length, Array.new(F32, 0n, 0.0)), valid}, exact}}}
342	
343	def built_owned(+kind: Pass, +input: Shape.Shape, +output: Shape.Shape, +plan: Plan, +depth: Nat, +length: Nat,
344	  -valid: {Storage.bounds(depth, length) == True{} : Bool},
345	  -exact: {length == Shape.logical_numel(Shape.dimensions(output)) : Nat},
346	  source: Core.Tensor<input>) -> Core.Dynamic & Result<&2, &1, Error, Core.Dynamic>:
347	  Core.TTensor{storage, fits} = source
348	  match storage:
349	    case Storage.Empty{}:
350	      (Core.Packed{input, Core.TTensor{Storage.Empty{}, fits}}, built_empty(kind, output, plan, depth, length, valid, exact))
351	    case Storage.Owned{+d, +n, array, balanced, ok}:
352	      built_cloned(kind, input, output, plan, depth, length, valid, exact, d, n, array, Array.clone(F32, array), {==}, balanced, ok, fits)
353	
354	# Every output cell at once, into the storage that `Core.new` allocates. An
355	# output with no cell keeps that storage, which is empty.
356	def built(+kind: Pass, +input: Shape.Shape, +output: Shape.Shape, +plan: Plan, source: Core.Tensor<input>,
357	  allocated: Result<&2, &1, Core.Error, Core.Tensor<output>>) -> Core.Dynamic & Result<&2, &1, Error, Core.Dynamic>:
358	  match allocated:
359	    case Fail{error}:
360	      (Core.Packed{input, source}, Fail{TensorError{error}})
361	    case Done{target}:
362	      Core.TTensor{storage, exact} = target
363	      match storage:
364	        case Storage.Empty{}:
365	          (Core.Packed{input, source}, Done{Core.Packed{output, Core.TTensor{Storage.Empty{}, exact}}})
366	        case Storage.Owned{+depth, +length, array, balanced, valid}:
367	          built_owned(kind, input, output, plan, depth, length, valid, exact, source)
368	
369	def sliced(+shape: Shape.Shape, tensor: Core.Tensor<shape>, planned: Result<&2, &2, Error, Planned>) ->
370	  Core.Dynamic & Result<&2, &1, Error, Core.Dynamic>:
371	  match planned:
372	    case Fail{error}:
373	      (Core.Packed{shape, tensor}, Fail{error})
374	    case Done{Planned{+plan, +output}}:
375	      built(Forward{}, shape, output, plan, tensor, Core.new(output, 0.0))
376	
377	# PyTorch's input[..., start:stop:step, ...] with the slice at axis `dim`.
378	# The input owner always comes back unchanged next to the result.
379	def slice(input: Core.Dynamic, +dim: U32, +start: U32, +stop: U32, +step: U32) -> Core.Dynamic & Result<&2, &1, Error, Core.Dynamic>:
380	  Core.Packed{+shape, tensor} = input
381	  sliced(shape, tensor, plan(shape, dim, start, stop, step))
382	
383	# PyTorch's slice_backward: the gradient g of a planned slice's output,
384	# scattered into +0.0 words of the input's shape. g must hold as many words
385	# as the slice's output; g comes back unchanged next to the result.
386	def scatter(g: Core.Dynamic, +input: Shape.Shape, +plan: Plan) -> Core.Dynamic & Result<&2, &1, Error, Core.Dynamic>:
387	  Core.Packed{+shape, tensor} = g
388	  built(Backward{}, shape, input, plan, tensor, Core.new(input, 0.0))
389	
390	def transformed(result: Result<&2, &1, Transforms.Error, Core.Dynamic>) -> Result<&2, &1, Error, Core.Dynamic>:
391	  match result:
392	    case Fail{Transforms.LayoutError{error}}:
393	      Fail{LayoutError{error}}
394	    case Fail{Transforms.TensorError{error}}:
395	      Fail{TensorError{error}}
396	    case Done{output}:
397	      Done{output}
398	
399	def reshaped(+dimensions: List<&2, U32>, pair: Core.Dynamic & Result<&2, &1, Error, Core.Dynamic>) ->
400	  Core.Dynamic & Result<&2, &1, Error, Core.Dynamic>:
401	  (input, result) = pair
402	  match result:
403	    case Fail{error}:
404	      (input, Fail{error})
405	    case Done{output}:
406	      (input, transformed(Transforms.reshape_dynamic(output, dimensions)))
407	
408	def narrow_ranged(+shape: Shape.Shape, tensor: Core.Tensor<shape>, +dim: U32, +start: U32, +length: U32, fits: Bool) ->
409	  Core.Dynamic & Result<&2, &1, Error, Core.Dynamic>:
410	  match fits:
411	    case False{}:
412	      (Core.Packed{shape, tensor}, Fail{OutOfRange{}})
413	    case True{}:
414	      slice(Core.Packed{shape, tensor}, dim, start, U32.add(start, length), 1)
415	
416	def narrow_checked(+shape: Shape.Shape, tensor: Core.Tensor<shape>, +dim: U32, +start: U32, +length: U32,
417	  refused: Maybe<&2, Error>) -> Core.Dynamic & Result<&2, &1, Error, Core.Dynamic>:
418	  match refused:
419	    case Some{error}:
420	      (Core.Packed{shape, tensor}, Fail{error})
421	    case None{}:
422	      narrow_ranged(shape, tensor, dim, start, length,
423	        Nat.is_le(Nat.add(U32.to_nat(start), U32.to_nat(length)), extent_at(Shape.dimensions(shape), dim)))
424	
425	# PyTorch's input.narrow(dim, start, length): the slice start:start + length
426	# at axis `dim`, refused when start + length passes the extent.
427	def narrow(input: Core.Dynamic, +dim: U32, +start: U32, +length: U32) -> Core.Dynamic & Result<&2, &1, Error, Core.Dynamic>:
428	  Core.Packed{+shape, tensor} = input
429	  narrow_checked(shape, tensor, dim, start, length, refusal(Shape.dimensions(shape), dim, 1))
430	
431	def select_ranged(+shape: Shape.Shape, tensor: Core.Tensor<shape>, +dim: U32, +index: U32, below: Bool) ->
432	  Core.Dynamic & Result<&2, &1, Error, Core.Dynamic>:
433	  match below:
434	    case False{}:
435	      (Core.Packed{shape, tensor}, Fail{OutOfRange{}})
436	    case True{}:
437	      reshaped(dropped(Shape.dimensions(shape), U32.to_nat(dim)), slice(Core.Packed{shape, tensor}, dim, index, U32.add(index, 1), 1))
438	
439	def select_checked(+shape: Shape.Shape, tensor: Core.Tensor<shape>, +dim: U32, +index: U32, refused: Maybe<&2, Error>) ->
440	  Core.Dynamic & Result<&2, &1, Error, Core.Dynamic>:
441	  match refused:
442	    case Some{error}:
443	      (Core.Packed{shape, tensor}, Fail{error})
444	    case None{}:
445	      select_ranged(shape, tensor, dim, index, Nat.is_lt(U32.to_nat(index), extent_at(Shape.dimensions(shape), dim)))
446	
447	# PyTorch's input.select(dim, index): the slice index:index + 1 at axis
448	# `dim`, then a reshape that drops the axis, refused when the index is at or
449	# past the extent.
450	def select(input: Core.Dynamic, +dim: U32, +index: U32) -> Core.Dynamic & Result<&2, &1, Error, Core.Dynamic>:
451	  Core.Packed{+shape, tensor} = input
452	  select_checked(shape, tensor, dim, index, refusal(Shape.dimensions(shape), dim, 1))
453	
454	def error_code(error: Error) -> U32:
455	  match error:
456	    case TensorError{error}:
457	      Core.error_code(error)
458	    case LayoutError{error}:
459	      Transforms.error_code(Transforms.LayoutError{error})
460	    case StepNotPositive{}:
461	      19
462	    case OutOfRange{}:
463	      20