src/tensor/slice.bend

bendtorch-core@0.1.0.0 · 463 lines · 60 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/transforms.bend as Transforms
10import bendtorch-core@0.1.0.0/src/tensor/parallel.bend as Parallel
11import bendtorch-core@0.1.0.0/src/tensor/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
41type 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)).
49type Bounds is Data:
50  Bounds{start: Nat, stop: Nat}
51
52def 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.
59def 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'.
68def 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.
74type View is Data:
75  View{outer: Nat, extent: Nat, inner: Nat}
76
77def 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
86def 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.
91def 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
100def 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.
104def 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.
115def 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.
123def 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.
133def 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.
144type Plan is Data:
145  Plan{outer: Nat, extent: Nat, inner: Nat, start: Nat, step: Nat, length: Nat}
146
147def 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
151def 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
155def 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
158def 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.
164def 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
173def 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.
177type Planned is Data:
178  Planned{plan: Plan, output: Shape.Shape}
179
180def 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
187def 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.
198def 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.
205def 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
211def same(+plan: Plan) -> Plan:
212  plan
213
214def take(b: Array<F32>, pair: Array<F32> & F32) -> F32 & (Array<F32> & Array<F32>):
215  (a, x) = pair
216  (x, (a, b))
217
218def 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.
224def 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.
232def splits(+depth: Nat) -> Nat:
233  Nat.sub(depth, 15n)
234
235def 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
239def 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.
250def 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
258def 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
261def 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
268def 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
272def 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
279def 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
286def 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
290def 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.
298type Pass is Data:
299  Forward{}
300  Backward{}
301
302def 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
309def 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.
319def 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.
333def 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
343def 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.
356def 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
369def 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.
379def 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.
386def 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
390def 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
399def 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
408def 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
416def 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.
427def 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
431def 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
439def 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.
450def 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
454def 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