src/tensor/signed.bend

bendtorch-core@0.1.0.0 · 366 lines · 48 definitions

1import Base
2import bendtorch-core@0.1.0.0/src/shape/core.bend as Shape
3import bendtorch-core@0.1.0.0/src/tensor/core.bend as Core
4import bendtorch-core@0.1.0.0/src/tensor/layout.bend as Layout
5import bendtorch-core@0.1.0.0/src/tensor/transforms.bend as Transforms
6import bendtorch-core@0.1.0.0/src/tensor/transpose_spec.bend as Transpose
7import bendtorch-core@0.1.0.0/src/tensor/slice.bend as Slice
8
9# PyTorch's signed axes, indices and reshape sizes, parsed into the unsigned
10# ones that the tensor operations check. A Position counts from either end, as
11# a PyTorch int does: First{k} is k, and Last{k} is -1 - k, so Last{0} is -1,
12# the last axis or element. Each integer has exactly one Position.
13type Position is Data:
14  First{k: U32}
15  Last{k: U32}
16
17# A reshape size: an exact size, or -1, the one size PyTorch infers.
18type Extent is Data:
19  Exact{size: U32}
20  Inferred{}
21
22# A signed reshape's refusals: the unsigned reshape's, more than one Inferred
23# size, and an Inferred size that any value would fit, on an empty input whose
24# other sizes multiply to 0.
25type Error is Data:
26  TransformError{error: Transforms.Error}
27  RepeatedInferred{}
28  AmbiguousInferred{}
29
30# PyTorch's wrap of a position among `limit` axes or elements (maybe_wrap_dim):
31# a position d in -limit..limit-1 becomes d + limit when d is negative and d
32# otherwise. A position outside that range stays outside it, at or past
33# `limit`, so the operation refuses it with its own check, code and order:
34# First{k} stays k, and an out-of-range Last{k} becomes `limit`.
35def wrap(+limit: U32, position: Position) -> U32:
36  match position:
37    case First{k}:
38      k
39    case Last{+k}:
40      Bool.pick(U32, U32.is_lt(k, limit), U32.sub(U32.sub(limit, 1), k), limit)
41
42def wrap_all(+limit: U32, positions: List<&2, Position>) -> List<&2, U32>:
43  match positions:
44    case Nil{}:
45      Nil{}
46    case head <> tail:
47      wrap(limit, head) <> wrap_all(limit, tail)
48
49# Each coordinate wrapped among its own axis's extent. A coordinate past the
50# rank has no extent and is wrapped among 0: the read refuses the coordinate
51# count before any range.
52def wrap_each(dimensions: List<&2, U32>, positions: List<&2, Position>) -> List<&2, U32>:
53  match dimensions positions:
54    case Nil{} positions:
55      wrap_all(0, positions)
56    case extent <> rest Nil{}:
57      Nil{}
58    case extent <> rest head <> tail:
59      wrap(extent, head) <> wrap_each(rest, tail)
60
61# Whether a position lies in PyTorch's range -limit..limit-1. Last{k} is
62# -1 - k, at least -limit exactly when k < limit.
63def inside(+limit: U32, position: Position) -> Bool:
64  match position:
65    case First{k}:
66      U32.is_lt(k, limit)
67    case Last{k}:
68      U32.is_lt(k, limit)
69
70def all_inside(+limit: U32, positions: List<&2, Position>) -> Bool:
71  match positions:
72    case Nil{}:
73      True{}
74    case head <> tail:
75      Bool.and(inside(limit, head), all_inside(limit, tail))
76
77# The axes of the input for transpose and sum: its rank, and 1 for a rank-0
78# input, whose only axis is 0 (or -1).
79def axis_limit(+input: Shape.Shape) -> U32:
80  U32.from_nat(Transpose.axis_limit(Shape.rank_dimensions(Shape.dimensions(input))))
81
82# How many sizes are Inferred.
83def inferred_count(extents: List<&2, Extent>) -> Nat:
84  match extents:
85    case Nil{}:
86      0n
87    case Exact{size} <> tail:
88      inferred_count(tail)
89    case Inferred{} <> tail:
90      1n+inferred_count(tail)
91
92# The product of the exact sizes, in Nat: what the laws state the inference
93# against. The inference never forms it, since it can pass the native Nat
94# limit of 2^48 on an empty input: [-1, 2^24, 2^24] multiplies to 2^48.
95def exact_product(extents: List<&2, Extent>) -> Nat:
96  match extents:
97    case Nil{}:
98      1n
99    case Exact{+size} <> tail:
100      Nat.mul(U32.to_nat(size), exact_product(tail))
101    case Inferred{} <> tail:
102      exact_product(tail)
103
104# The sizes, each Inferred one replaced by `size`.
105def filled(extents: List<&2, Extent>, +size: U32) -> List<&2, U32>:
106  match extents:
107    case Nil{}:
108      Nil{}
109    case Exact{extent} <> tail:
110      extent <> filled(tail, size)
111    case Inferred{} <> tail:
112      size <> filled(tail, size)
113
114# The exact product against a bound, the input's element count: ZeroProduct
115# when it is 0, Within{p} when it is p > 0 at most the bound, and Beyond when
116# it is past the bound. capped only multiplies a size and a product that are
117# both at most the bound, and a tensor holds at most 2^20 elements, so no
118# number it forms passes 2^40.
119type Product is Data:
120  ZeroProduct{}
121  Within{product: Nat}
122  Beyond{}
123
124def bounded(product: Nat, below: Bool) -> Product:
125  match below:
126    case True{}:
127      Within{product}
128    case False{}:
129      Beyond{}
130
131def cap(+bound: Nat, +product: Nat) -> Product:
132  bounded(product, Nat.is_le(product, bound))
133
134def scaled(+bound: Nat, +size: Nat, +product: Nat, past: Bool) -> Product:
135  match past:
136    case True{}:
137      Beyond{}
138    case False{}:
139      cap(bound, Nat.mul(size, product))
140
141# A nonzero size times the product of the sizes after it. A size past the
142# bound makes any nonzero product pass it, so it is never multiplied.
143def times(+bound: Nat, size: Nat, rest: Product) -> Product:
144  match size:
145    case 0n:
146      ZeroProduct{}
147    case 1n+(+less):
148      match rest:
149        case ZeroProduct{}:
150          ZeroProduct{}
151        case Within{+product}:
152          scaled(bound, 1n+less, product, Nat.is_lt(bound, 1n+less))
153        case Beyond{}:
154          Beyond{}
155
156def capped(+bound: Nat, extents: List<&2, Extent>) -> Product:
157  match extents:
158    case Nil{}:
159      cap(bound, 1n)
160    case Exact{+size} <> tail:
161      times(bound, U32.to_nat(size), capped(bound, tail))
162    case Inferred{} <> tail:
163      capped(bound, tail)
164
165def count_mismatch() -> Error:
166  TransformError{Transforms.LayoutError{Layout.ElementCountMismatch{}}}
167
168# An Inferred size whose other sizes multiply to 0: ambiguous on an empty
169# input, a count mismatch on any other.
170def unfilled(count: Nat) -> Result<&2, &2, Error, List<&2, U32>>:
171  match count:
172    case 0n:
173      Fail{AmbiguousInferred{}}
174    case 1n+rest:
175      Fail{count_mismatch()}
176
177def divided(+extents: List<&2, Extent>, split: Nat & Nat) -> Result<&2, &2, Error, List<&2, U32>>:
178  (quotient, remainder) = split
179  match remainder:
180    case 0n:
181      Done{filled(extents, U32.from_nat(quotient))}
182    case 1n+rest:
183      Fail{count_mismatch()}
184
185# A product past the count: the inferred size is 0 on an empty input, and no
186# size fits any other.
187def beyond(count: Nat, +extents: List<&2, Extent>) -> Result<&2, &2, Error, List<&2, U32>>:
188  match count:
189    case 0n:
190      Done{filled(extents, 0)}
191    case 1n+rest:
192      Fail{count_mismatch()}
193
194def inferred_one(+count: Nat, +extents: List<&2, Extent>, product: Product) -> Result<&2, &2, Error, List<&2, U32>>:
195  match product:
196    case ZeroProduct{}:
197      unfilled(count)
198    case Within{product}:
199      divided(extents, Nat.divmod(count, product))
200    case Beyond{}:
201      beyond(count, extents)
202
203def inferred_from(+count: Nat, +extents: List<&2, Extent>, inferred: Nat) -> Result<&2, &2, Error, List<&2, U32>>:
204  match inferred:
205    case 0n:
206      Done{filled(extents, 0)}
207    case 1n+more:
208      match more:
209        case 0n:
210          inferred_one(count, extents, capped(count, extents))
211        case 1n+rest:
212          Fail{RepeatedInferred{}}
213
214# PyTorch's infer_size: the sizes of a reshape of `count` elements. With no
215# Inferred size, the sizes themselves; with one, it is count / product of
216# the others when that divides exactly, the product capped at the count;
217# more than one is refused first.
218def inferred_dimensions(+count: Nat, +extents: List<&2, Extent>) -> Result<&2, &2, Error, List<&2, U32>>:
219  inferred_from(count, extents, inferred_count(extents))
220
221def transformed(result: Result<&2, &1, Transforms.Error, Core.Dynamic>) -> Result<&2, &1, Error, Core.Dynamic>:
222  match result:
223    case Fail{error}:
224      Fail{TransformError{error}}
225    case Done{output}:
226      Done{output}
227
228def reshape_inferred(input: Shape.Shape, tensor: Core.Tensor<input>,
229  inferred: Result<&2, &2, Error, List<&2, U32>>) -> Result<&2, &1, Error, Core.Dynamic>:
230  match inferred:
231    case Fail{error}:
232      Fail{error}
233    case Done{+dimensions}:
234      transformed(Transforms.reshape(input, tensor, dimensions))
235
236# PyTorch's reshape with -1: the sizes inferred, then the unsigned reshape.
237def reshape(+input: Shape.Shape, tensor: Core.Tensor<input>, +extents: List<&2, Extent>) -> Result<&2, &1, Error, Core.Dynamic>:
238  reshape_inferred(input, tensor, inferred_dimensions(Shape.elements(input), extents))
239
240def reshape_dynamic(dynamic: Core.Dynamic, extents: List<&2, Extent>) -> Result<&2, &1, Error, Core.Dynamic>:
241  Core.Packed{+input, tensor} = dynamic
242  reshape(input, tensor, extents)
243
244# PyTorch's permute_copy with signed axes, wrapped among the input's rank.
245def permute_copy(+input: Shape.Shape, source: Core.Tensor<input>, axes: List<&2, Position>) ->
246  Core.Tensor<input> & Result<&2, &1, Transforms.Error, Core.Dynamic>:
247  Transforms.permute_copy(input, source, wrap_all(Shape.rank(input), axes))
248
249def permute_dynamic(dynamic: Core.Dynamic, axes: List<&2, Position>) -> Core.Dynamic & Result<&2, &1, Transforms.Error, Core.Dynamic>:
250  Core.Packed{+input, source} = dynamic
251  Transforms.dynamic_result(input, permute_copy(input, source, axes))
252
253# PyTorch's transpose_copy with signed axes, wrapped among axis_limit.
254def transpose_copy(+input: Shape.Shape, source: Core.Tensor<input>, a: Position, b: Position) ->
255  Core.Tensor<input> & Result<&2, &1, Transforms.Error, Core.Dynamic>:
256  Transforms.transpose_copy(input, source, wrap(axis_limit(input), a), wrap(axis_limit(input), b))
257
258def transpose_dynamic(dynamic: Core.Dynamic, a: Position, b: Position) -> Core.Dynamic & Result<&2, &1, Transforms.Error, Core.Dynamic>:
259  Core.Packed{+input, source} = dynamic
260  Transforms.dynamic_result(input, transpose_copy(input, source, a, b))
261
262# A[coordinates] with signed coordinates, each wrapped among its axis's
263# extent; and the write A[coordinates] = value.
264def read_coordinate(+shape: Shape.Shape, tensor: Core.Tensor<shape>, coordinates: List<&2, Position>) ->
265  Core.Tensor<shape> & Result<&2, &2, Core.Error, F32>:
266  Core.read_coordinate(shape, tensor, wrap_each(Shape.dimensions(shape), coordinates))
267
268def write_coordinate(+shape: Shape.Shape, tensor: Core.Tensor<shape>, coordinates: List<&2, Position>, value: F32) ->
269  Core.Tensor<shape> & Result<&2, &2, Core.Error, Unit>:
270  Core.write_coordinate(shape, tensor, wrap_each(Shape.dimensions(shape), coordinates), value)
271
272# The element at a signed flat index, PyTorch's A.flatten()[index], wrapped
273# among the element count; and the write at it.
274def read_flat(+shape: Shape.Shape, tensor: Core.Tensor<shape>, index: Position) ->
275  Core.Tensor<shape> & Result<&2, &2, Core.Error, F32>:
276  Core.read_flat(shape, tensor, wrap(Shape.count(shape), index))
277
278def write_flat(+shape: Shape.Shape, tensor: Core.Tensor<shape>, index: Position, +value: F32) ->
279  Core.Tensor<shape> & Result<&2, &2, Core.Error, Unit>:
280  Core.write_flat(shape, tensor, wrap(Shape.count(shape), index), value)
281
282def read_dynamic(dynamic: Core.Dynamic, index: Position) -> Core.Dynamic & Result<&2, &2, Core.Error, F32>:
283  Core.Packed{+shape, tensor} = dynamic
284  Core.read_dynamic_result(shape, read_flat(shape, tensor, index))
285
286def write_dynamic(dynamic: Core.Dynamic, index: Position, value: F32) -> Core.Dynamic & Result<&2, &2, Core.Error, Unit>:
287  Core.Packed{+shape, tensor} = dynamic
288  Core.write_dynamic_result(shape, write_flat(shape, tensor, index, value))
289
290# Whether a position lies before -limit: Last{k}, the position -1 - k, with
291# k at or past limit. No First does.
292def before(+limit: U32, position: Position) -> Bool:
293  match position:
294    case First{k}:
295      False{}
296    case Last{k}:
297      Bool.not(U32.is_lt(k, limit))
298
299# PyTorch's slice positions (aten::slice, get_slice_bounds) before its clamp
300# to the axis: a negative position d becomes d + limit, the wrap, and one
301# still below 0 becomes 0. First{k} stays k; Slice.bounds then clamps both
302# ends to the extent.
303def bound(+limit: U32, +position: Position) -> U32:
304  Bool.pick(U32, before(limit, position), 0, wrap(limit, position))
305
306# The extent of axis `axis`, 0 past the rank.
307def extent(dimensions: List<&2, U32>, +axis: U32) -> U32:
308  U32.from_nat(Slice.extent_at(dimensions, axis))
309
310# PyTorch's input[..., start:stop:step, ...] with the slice at a signed axis
311# `dim`, wrapped among the rank, and signed start and stop: the unsigned
312# slice of the bounds. A rank-0 input has no axis, so every dim is refused.
313def slice(input: Core.Dynamic, dim: Position, start: Position, stop: Position, +step: U32) ->
314  Core.Dynamic & Result<&2, &1, Slice.Error, Core.Dynamic>:
315  Core.Packed{+shape, tensor} = input
316  +axis = wrap(Shape.rank(shape), dim)
317  +n = extent(Shape.dimensions(shape), axis)
318  Slice.slice(Core.Packed{shape, tensor}, axis, bound(n, start), bound(n, stop), step)
319
320# PyTorch's input.select(dim, index) at a signed axis and a signed index,
321# wrapped among the extent: an index outside -extent..extent-1 stays at or
322# past the extent, so the unsigned select refuses it with OutOfRange after
323# its axis check.
324def select(input: Core.Dynamic, dim: Position, index: Position) -> Core.Dynamic & Result<&2, &1, Slice.Error, Core.Dynamic>:
325  Core.Packed{+shape, tensor} = input
326  +axis = wrap(Shape.rank(shape), dim)
327  Slice.select(Core.Packed{shape, tensor}, axis, wrap(extent(Shape.dimensions(shape), axis), index))
328
329# A narrow's start before -extent: the axis check first, as the unsigned
330# narrow does it, then OutOfRange.
331def narrow_before(input: Core.Dynamic, refused: Maybe<&2, Slice.Error>) -> Core.Dynamic & Result<&2, &1, Slice.Error, Core.Dynamic>:
332  match refused:
333    case Some{error}:
334      (input, Fail{error})
335    case None{}:
336      (input, Fail{Slice.OutOfRange{}})
337
338def narrow_started(+shape: Shape.Shape, tensor: Core.Tensor<shape>, +axis: U32, +start: U32, +length: U32, early: Bool) ->
339  Core.Dynamic & Result<&2, &1, Slice.Error, Core.Dynamic>:
340  match early:
341    case False{}:
342      Slice.narrow(Core.Packed{shape, tensor}, axis, start, length)
343    case True{}:
344      narrow_before(Core.Packed{shape, tensor}, Slice.refusal(Shape.dimensions(shape), axis, 1))
345
346# PyTorch's input.narrow(dim, start, length) at a signed axis with a signed
347# start. PyTorch takes a start in -extent..extent: Last{k} with k below the
348# extent is extent - 1 - k, and one before -extent is OutOfRange. A First
349# start past the extent is refused by the unsigned narrow, which checks
350# start + length against the extent, with the same code. A length is a
351# U32, so PyTorch's negative lengths cannot be written.
352def narrow(input: Core.Dynamic, dim: Position, +start: Position, +length: U32) ->
353  Core.Dynamic & Result<&2, &1, Slice.Error, Core.Dynamic>:
354  Core.Packed{+shape, tensor} = input
355  +axis = wrap(Shape.rank(shape), dim)
356  +n = extent(Shape.dimensions(shape), axis)
357  narrow_started(shape, tensor, axis, wrap(n, start), length, before(n, start))
358
359def error_code(error: Error) -> U32:
360  match error:
361    case TransformError{error}:
362      Transforms.error_code(error)
363    case RepeatedInferred{}:
364      17
365    case AmbiguousInferred{}:
366      18