src/tensor/elementwise.bend

bendtorch-core@0.1.0.0 · 665 lines · 86 definitions

1import Base
2import bendtorch-core@0.1.0.0/src/shape/core.bend as Shape
3import bendtorch-core@0.1.0.0/src/shape/flat.bend as Flat
4import bendtorch-core@0.1.0.0/src/storage/main.bend as Storage
5import bendtorch-core@0.1.0.0/src/storage/model.bend as Model
6import bendtorch-core@0.1.0.0/src/storage/certificates.bend as Certificates
7import bendtorch-core@0.1.0.0/src/tensor/core.bend as Core
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
10
11type Operation is Data:
12  Add{}
13  Sub{}
14  Mul{}
15  KeepLeft{}
16  NegateLeft{}
17  ReluLeft{}
18  ReluBackward{}
19  Divide{}
20  ExpLeft{}
21  LogLeft{}
22  Matches{}
23  PickNegated{}
24  MseGrad{scale: F32}
25  OneHotGrad{count: F32}
26  SoftmaxGrad{count: F32, single: Bool}
27  Descend{lr: F32}
28
29# The IEEE 754 bits of an F32, as a U32.
30def bits(x: F32) -> U32:
31  F32{word} = x
32  U32{word}
33
34# IEEE 754 `x < 0`, on the bits: the words above -0.0 (0x80000000) up to
35# -inf (0xFF800000). -0.0, +0.0 and every NaN are not below zero.
36def below_zero(x: F32) -> Bool:
37  +word = bits(x)
38  Bool.and(U32.is_lt(2147483648, word), U32.is_le(word, 4286578688))
39
40# IEEE 754 `x <= 0`, on the bits: +0.0 (0x00000000) and the words from -0.0
41# (0x80000000) up to -inf (0xFF800000). No NaN is at most zero.
42def at_most_zero(x: F32) -> Bool:
43  +word = bits(x)
44  Bool.or(U32.is_eq(word, 0), Bool.and(U32.is_le(2147483648, word), U32.is_le(word, 4286578688)))
45
46# +0.0 when `zeroed`, else the word itself.
47def zero_if(zeroed: Bool, value: F32) -> F32:
48  match zeroed:
49    case True{}:
50      0.0
51    case False{}:
52      value
53
54# PyTorch's relu, clamp_min(x, 0): +0.0 where x < 0, else x's own word, so
55# -0.0 stays -0.0 and a NaN keeps its payload.
56def relu(+x: F32) -> F32:
57  zero_if(below_zero(x), x)
58
59# PyTorch's relu gradient, threshold_backward(g, relu(x), 0): +0.0 where
60# x <= 0, else g's own word. At x = +0.0 and -0.0 it is +0.0; at a NaN x it
61# is g.
62def relu_backward(+x: F32, g: F32) -> F32:
63  zero_if(at_most_zero(x), g)
64
65# Whether the word is 1.0 (0x3F800000), the target mark of a one-hot row.
66def is_one(x: F32) -> Bool:
67  U32.is_eq(bits(x), 1065353216)
68
69# -(g / count), the gradient PyTorch's nll_loss with mean reduction gives the
70# target cell of each row.
71def target_gradient(g: F32, +count: F32) -> F32:
72  F32.neg(F32.div(g, count))
73
74# The row sum PyTorch's log_softmax backward takes of a one-hot gradient row
75# whose target word is t. A row of one column holds t alone, and PyTorch's
76# sum of one word is that word, -0.0 kept; a wider row adds its +0.0 words,
77# so the sum is t + 0.0.
78def row_sum(single: Bool, t: F32) -> F32:
79  match single:
80    case True{}:
81      t
82    case False{}:
83      F32.add(t, 0.0)
84
85def scalar(operation: Operation, left: F32, right: F32) -> F32:
86  match operation:
87    case Add{}:
88      F32.add(left, right)
89    case Sub{}:
90      F32.sub(left, right)
91    case Mul{}:
92      F32.mul(left, right)
93    case KeepLeft{}:
94      left
95    case NegateLeft{}:
96      F32.neg(left)
97    case ReluLeft{}:
98      relu(left)
99    case ReluBackward{}:
100      relu_backward(left, right)
101    case Divide{}:
102      F32.div(left, right)
103    case ExpLeft{}:
104      F32.exp(left)
105    case LogLeft{}:
106      F32.log(left)
107    case Matches{}:
108      zero_if(Bool.not(U32.is_eq(bits(left), bits(right))), 1.0)
109    case PickNegated{}:
110      zero_if(Bool.not(is_one(right)), F32.neg(left))
111    case MseGrad{+scale}:
112      F32.mul(F32.mul(scale, right), left)
113    case OneHotGrad{+count}:
114      zero_if(Bool.not(is_one(right)), target_gradient(left, count))
115    case SoftmaxGrad{+count, +single}:
116      F32.neg(F32.mul(F32.exp(right), row_sum(single, target_gradient(left, count))))
117    case Descend{lr}:
118      F32.sub(left, F32.mul(lr, right))
119
120def flat_input_indices(+plan: Shape.BroadcastPlan,
121  result: Result<&2, &2, Shape.Error, Shape.FlatIndex<Shape.output_shape(plan)>>) -> Result<&2, &1, Core.Error, U32 & U32>:
122  match result:
123    case Fail{error}:
124      Fail{Core.ShapeError{error}}
125    case Done{+index}:
126      +map = Shape.index_map(plan)
127      Done{(Shape.index_value(Shape.left_shape(plan), Flat.project_left_index(plan, map, index)),
128        Shape.index_value(Shape.right_shape(plan), Flat.project_right_index(plan, map, index)))}
129
130def input_indices(+plan: Shape.BroadcastPlan, index: U32) -> Result<&2, &1, Core.Error, U32 & U32>:
131  flat_input_indices(plan, Shape.flat_checked(Shape.output_shape(plan), U32.to_nat(index)))
132
133type State<-shape: Shape.Shape> is Type:
134  State{left: Core.Dynamic, right: Core.Dynamic, output: Core.Tensor<shape>}
135
136type Step<-shape: Shape.Shape> is Type:
137  Continue{state: State<shape>}
138  Stop{left: Core.Dynamic, right: Core.Dynamic, error: Core.Error}
139
140def step_write(-shape: Shape.Shape, left: Core.Dynamic, right: Core.Dynamic,
141  pair: Core.Tensor<shape> & Result<&2, &2, Core.Error, Unit>) -> Step<shape>:
142  (output, result) = pair
143  match result:
144    case Fail{error}:
145      Stop{left, right, error}
146    case Done{unit}:
147      Continue{State{left, right, output}}
148
149def step_right(+operation: Operation, -shape: Shape.Shape, +index: U32, left: Core.Dynamic,
150  output: Core.Tensor<shape>, left_value: F32, pair: Core.Dynamic & Result<&2, &2, Core.Error, F32>) -> Step<shape>:
151  (right, result) = pair
152  match result:
153    case Fail{error}:
154      Stop{left, right, error}
155    case Done{right_value}:
156      step_write(shape, left, right, Core.write_flat(shape, output, index, scalar(operation, left_value, right_value)))
157
158def step_left(+operation: Operation, -shape: Shape.Shape, +index: U32, right: Core.Dynamic,
159  output: Core.Tensor<shape>, right_index: U32, pair: Core.Dynamic & Result<&2, &2, Core.Error, F32>) -> Step<shape>:
160  (left, result) = pair
161  match result:
162    case Fail{error}:
163      Stop{left, right, error}
164    case Done{left_value}:
165      step_right(operation, shape, index, left, output, left_value, Core.read_dynamic(right, right_index))
166
167def step_indices(+operation: Operation, -shape: Shape.Shape, +index: U32, state: State<shape>,
168  result: Result<&2, &1, Core.Error, U32 & U32>) -> Step<shape>:
169  State{left, right, output} = state
170  match result:
171    case Fail{error}:
172      Stop{left, right, error}
173    case Done{indices}:
174      (left_index, right_index) = indices
175      step_left(operation, shape, index, right, output, right_index, Core.read_dynamic(left, left_index))
176
177def step(+operation: Operation, +plan: Shape.BroadcastPlan, +index: U32, state: State<Shape.output_shape(plan)>) ->
178  Step<Shape.output_shape(plan)>:
179  step_indices(operation, Shape.output_shape(plan), index, state, input_indices(plan, index))
180
181type NaturalStep<-shape: Shape.Shape, -cursor: Nat> is Type:
182  NaturalContinue{state: State<shape>, -index: U32, -exact: {U32.to_nat(index) == cursor : Nat}}
183  NaturalStop{left: Core.Dynamic, right: Core.Dynamic, error: Core.Error}
184
185type CursorWitness<-cursor: Nat, -index: U32> is Data:
186  CursorWitness{-exact: {U32.to_nat(index) == cursor : Nat}}
187
188def CursorValidity(-shape: Shape.Shape, +cursor: Nat, result: NaturalStep<shape, cursor>) -> Data:
189  match result:
190    case NaturalStop{left, right, error}:
191      Unit
192    case NaturalContinue{state, index, exact}:
193      CursorWitness<cursor, index>
194
195def natural_step_result(-shape: Shape.Shape, -cursor: Nat, -index: U32,
196  -exact: {U32.to_nat(index) == cursor : Nat}, result: Step<shape>) -> NaturalStep<shape, cursor>:
197  match result:
198    case Stop{left, right, error}:
199      NaturalStop{left, right, error}
200    case Continue{state}:
201      NaturalContinue{state, index, exact}
202
203def natural_conversion_error(-shape: Shape.Shape, -cursor: Nat, state: State<shape>) -> NaturalStep<shape, cursor>:
204  State{left, right, output} = state
205  NaturalStop{left, right, Core.StorageError{Storage.InvalidIndex{}}}
206
207def natural_step_decided(+operation: Operation, +plan: Shape.BroadcastPlan, +cursor: Nat, +index: U32,
208  state: State<Shape.output_shape(plan)>, decision: Shape.BoolCase<Nat.is_eq(U32.to_nat(index), cursor)>) ->
209  NaturalStep<Shape.output_shape(plan), cursor>:
210  match decision:
211    case Shape.BFalse{equal}:
212      natural_conversion_error(Shape.output_shape(plan), cursor, state)
213    case Shape.BTrue{equal}:
214      natural_step_result(Shape.output_shape(plan), cursor, index,
215        Lifecycle.natural_equal(U32.to_nat(index), cursor, equal), step(operation, plan, index, state))
216
217def natural_step(+operation: Operation, +plan: Shape.BroadcastPlan, +cursor: Nat, state: State<Shape.output_shape(plan)>) ->
218  NaturalStep<Shape.output_shape(plan), cursor>:
219  +index = U32.from_nat(cursor)
220  natural_step_decided(operation, plan, cursor, index, state, Shape.bool_case(Nat.is_eq(U32.to_nat(index), cursor)))
221
222def finish_step(-shape: Shape.Shape, -cursor: Nat,
223  rec: State<shape> -> (Core.Dynamic & Core.Dynamic) & Result<&2, &1, Core.Error, Core.Dynamic>,
224  result: NaturalStep<shape, cursor>) -> (Core.Dynamic & Core.Dynamic) & Result<&2, &1, Core.Error, Core.Dynamic>:
225  match result:
226    case NaturalStop{left, right, error}:
227      ((left, right), Fail{error})
228    case NaturalContinue{state, index, exact}:
229      rec(state)
230
231def loop(fuel: Nat, +operation: Operation, +plan: Shape.BroadcastPlan, +cursor: Nat, state: State<Shape.output_shape(plan)>) ->
232  (Core.Dynamic & Core.Dynamic) & Result<&2, &1, Core.Error, Core.Dynamic>:
233  match fuel:
234    case 0n:
235      State{left, right, output} = state
236      ((left, right), Done{Core.pack(Shape.output_shape(plan), output)})
237    case 1n+remaining:
238      finish_step(Shape.output_shape(plan), cursor,
239        state => loop(remaining, operation, plan, 1n+cursor, state), natural_step(operation, plan, cursor, state))
240
241def allocated(+operation: Operation, +plan: Shape.BroadcastPlan, left: Core.Dynamic, right: Core.Dynamic,
242  result: Result<&2, &1, Core.Error, Core.Tensor<Shape.output_shape(plan)>>) ->
243  (Core.Dynamic & Core.Dynamic) & Result<&2, &1, Core.Error, Core.Dynamic>:
244  match result:
245    case Fail{error}:
246      ((left, right), Fail{error})
247    case Done{output}:
248      loop(Shape.elements(Shape.output_shape(plan)), operation, plan, 0n, State{left, right, output})
249
250def left_walk(-left: List<&2, Nat>, -right: List<&2, Nat>, -output: List<&2, Nat>,
251  map: Shape.Index<left, right, output>, +position: Nat) -> Nat:
252  match map:
253    case Shape.Indexed{rows, exact}:
254      Shape.walk_left(rows, position)
255    case Shape.Unindexed{zero}:
256      0n
257
258def right_walk(-left: List<&2, Nat>, -right: List<&2, Nat>, -output: List<&2, Nat>,
259  map: Shape.Index<left, right, output>, +position: Nat) -> Nat:
260  match map:
261    case Shape.Indexed{rows, exact}:
262      Shape.walk_right(rows, position)
263    case Shape.Unindexed{zero}:
264      0n
265
266# Whether the left input's stride is the output's on every axis and the last
267# one is 1, so that the walk of every position is the position itself. Equal
268# left and output shapes with no axis of extent 1 have such rows.
269def same_left(rows: List<&2, Shape.Row>) -> Bool:
270  match rows:
271    case Shape.Row{+stride, +left, right} <> Nil{}:
272      Bool.and(Nat.is_eq(stride, 1n), Nat.is_eq(left, 1n))
273    case Shape.Row{+stride, +left, right} <> rest:
274      Bool.and(Nat.is_eq(left, stride), same_left(rest))
275    case Nil{}:
276      False{}
277
278def same_right(rows: List<&2, Shape.Row>) -> Bool:
279  match rows:
280    case Shape.Row{+stride, left, +right} <> Nil{}:
281      Bool.and(Nat.is_eq(stride, 1n), Nat.is_eq(right, 1n))
282    case Shape.Row{+stride, left, +right} <> rest:
283      Bool.and(Nat.is_eq(right, stride), same_right(rest))
284    case Nil{}:
285      False{}
286
287# Whether the left input's stride is 0 on every axis, as for a rank-0
288# input, so that every position walks to slot 0.
289def zero_left(rows: List<&2, Shape.Row>) -> Bool:
290  match rows:
291    case Shape.Row{stride, +left, right} <> rest:
292      Bool.and(Nat.is_eq(left, 0n), zero_left(rest))
293    case Nil{}:
294      True{}
295
296def zero_right(rows: List<&2, Shape.Row>) -> Bool:
297  match rows:
298    case Shape.Row{stride, left, +right} <> rest:
299      Bool.and(Nat.is_eq(right, 0n), zero_right(rest))
300    case Nil{}:
301      True{}
302
303# Where an input reads output cell p: at its own slot p (`OwnSlot`), always
304# at slot 0 (`FirstSlot`), or wherever the broadcast walk of p lands
305# (`WalkedSlot`).
306type Slot is Data:
307  OwnSlot{}
308  FirstSlot{}
309  WalkedSlot{}
310
311def slot_of(same: Bool, zero: Bool) -> Slot:
312  match same:
313    case True{}:
314      OwnSlot{}
315    case False{}:
316      match zero:
317        case True{}:
318          FirstSlot{}
319        case False{}:
320          WalkedSlot{}
321
322def left_slot(-left: List<&2, Nat>, -right: List<&2, Nat>, -output: List<&2, Nat>, map: Shape.Index<left, right, output>) -> Slot:
323  match map:
324    case Shape.Indexed{+rows, exact}:
325      slot_of(same_left(rows), zero_left(rows))
326    case Shape.Unindexed{zero}:
327      WalkedSlot{}
328
329def right_slot(-left: List<&2, Nat>, -right: List<&2, Nat>, -output: List<&2, Nat>, map: Shape.Index<left, right, output>) -> Slot:
330  match map:
331    case Shape.Indexed{+rows, exact}:
332      slot_of(same_right(rows), zero_right(rows))
333    case Shape.Unindexed{zero}:
334      WalkedSlot{}
335
336# One output cell: `scalar(operation, left[ia], right[ib])` when the cell is
337# below `count`, else `+0.0`. Both input arrays come back with the cell.
338def cell_right(+operation: Operation, x: F32, la: Array<F32>, pair: Array<F32> & F32) ->
339  F32 & (Array<F32> & Array<F32>):
340  (ra, y) = pair
341  (scalar(operation, x, y), (la, ra))
342
343def cell_left(+operation: Operation, ra: Array<F32>, +right_index: U32, pair: Array<F32> & F32) ->
344  F32 & (Array<F32> & Array<F32>):
345  (la, x) = pair
346  cell_right(operation, x, la, Array.get(F32, ra, right_index))
347
348def cell(+operation: Operation, -le: List<&2, Nat>, -re: List<&2, Nat>, -oe: List<&2, Nat>,
349  +map: Shape.Index<le, re, oe>, +position: Nat, la: Array<F32>, ra: Array<F32>, inside: Bool) ->
350  F32 & (Array<F32> & Array<F32>):
351  match inside:
352    case False{}:
353      (0.0, (la, ra))
354    case True{}:
355      cell_left(operation, ra, U32.from_nat(right_walk(le, re, oe, map, position)),
356        Array.get(F32, la, U32.from_nat(left_walk(le, re, oe, map, position))))
357
358# A build's operation and broadcast plan.
359type Request is Data:
360  Request{operation: Operation, plan: Shape.BroadcastPlan}
361
362# What every cell of a build needs: the operation, the broadcast index map
363# and the slot each input reads (`left_slot`, `right_slot`), so that a build
364# whose inputs have the output's shape or rank 0 skips the walk. Each thread
365# below the forks computes its own once.
366type Walk is Data:
367  Walk{operation: Operation, -le: List<&2, Nat>, -re: List<&2, Nat>, -oe: List<&2, Nat>, map: Shape.Index<le, re, oe>,
368    left: Slot, right: Slot}
369
370def walk(+request: Request) -> Walk:
371  Request{+operation, +plan} = request
372  +le = Shape.extents(Shape.left_shape(plan))
373  +re = Shape.extents(Shape.right_shape(plan))
374  +oe = Shape.extents(Shape.output_shape(plan))
375  Walk{operation, le, re, oe, Shape.index_map(plan), left_slot(le, re, oe, Shape.index_map(plan)),
376    right_slot(le, re, oe, Shape.index_map(plan))}
377
378# An inside cell whose input slots are already known: the cell `cell`
379# computes there.
380def direct_cell(+operation: Operation, +left_index: Nat, +right_index: Nat, la: Array<F32>, ra: Array<F32>) ->
381  F32 & (Array<F32> & Array<F32>):
382  cell_left(operation, ra, U32.from_nat(right_index), Array.get(F32, la, U32.from_nat(left_index)))
383
384# A padding cell is the one `cell` computes; an inside cell reads each input
385# at the slot its `Slot` names.
386def cell_by(inside: Bool, +left: Slot, +right: Slot, +operation: Operation, -le: List<&2, Nat>, -re: List<&2, Nat>, -oe: List<&2, Nat>,
387  +map: Shape.Index<le, re, oe>, +position: Nat, la: Array<F32>, ra: Array<F32>) ->
388  F32 & (Array<F32> & Array<F32>):
389  match inside left right:
390    case False{} left right:
391      cell(operation, le, re, oe, map, position, la, ra, False{})
392    case True{} OwnSlot{} OwnSlot{}:
393      direct_cell(operation, position, position, la, ra)
394    case True{} OwnSlot{} FirstSlot{}:
395      direct_cell(operation, position, 0n, la, ra)
396    case True{} OwnSlot{} WalkedSlot{}:
397      direct_cell(operation, position, right_walk(le, re, oe, map, position), la, ra)
398    case True{} FirstSlot{} OwnSlot{}:
399      direct_cell(operation, 0n, position, la, ra)
400    case True{} FirstSlot{} FirstSlot{}:
401      direct_cell(operation, 0n, 0n, la, ra)
402    case True{} FirstSlot{} WalkedSlot{}:
403      direct_cell(operation, 0n, right_walk(le, re, oe, map, position), la, ra)
404    case True{} WalkedSlot{} OwnSlot{}:
405      direct_cell(operation, left_walk(le, re, oe, map, position), position, la, ra)
406    case True{} WalkedSlot{} FirstSlot{}:
407      direct_cell(operation, left_walk(le, re, oe, map, position), 0n, la, ra)
408    case True{} WalkedSlot{} WalkedSlot{}:
409      cell(operation, le, re, oe, map, position, la, ra, True{})
410
411# The cell function `Parallel.build` runs: its state is the pair of input
412# arrays.
413def cell_at(+walk: Walk, +position: Nat, inputs: Array<F32> & Array<F32>, inside: Bool) -> F32 & (Array<F32> & Array<F32>):
414  Walk{+operation, le, re, oe, +map, +left, +right} = walk
415  (la, ra) = inputs
416  cell_by(inside, left, right, operation, le, re, oe, map, position, la, ra)
417
418# Each half of a fork gets its own clone of both input arrays, because an
419# Array has one owner.
420def halves(lp: Array<F32> & Array<F32>, rp: Array<F32> & Array<F32>) -> (Array<F32> & Array<F32>) & (Array<F32> & Array<F32>):
421  (l1, l2) = lp
422  (r1, r2) = rp
423  ((l1, r1), (l2, r2))
424
425def split(inputs: Array<F32> & Array<F32>) -> (Array<F32> & Array<F32>) & (Array<F32> & Array<F32>):
426  (la, ra) = inputs
427  halves(Array.clone(F32, la), Array.clone(F32, ra))
428
429# The parallel build, `Parallel.build` with the elementwise cell: the top
430# `splits` levels fork, and below the forks a thread builds its subtree in
431# index order. Every cell gets the F32 term the sequential loop writes
432# there, so no output word depends on the split.
433def build(+depth: Nat, +splits: Nat, +operation: Operation, +plan: Shape.BroadcastPlan, +count: Nat, +base: Nat,
434  la: Array<F32>, ra: Array<F32>) -> Array<F32>:
435  Parallel.build(~Request, ~Walk, ~Array<F32>, ~Array<F32>, ~walk, ~split, ~cell_at, Request{operation, plan}, depth, splits, count, base, (la, ra))
436
437# A subtree below the forks holds 2^15 cells: an output of 2^20 cells forks
438# five levels into 32 subtrees, and an output of at most 2^15 cells does not
439# fork.
440def splits(+depth: Nat) -> Nat:
441  Nat.sub(depth, 15n)
442
443# The parallel build returns a balanced tree of its depth, whatever the
444# split: the erased certificate of the output storage.
445def build_balanced(+depth: Nat, +splits: Nat, +operation: Operation, +plan: Shape.BroadcastPlan, +count: Nat, +base: Nat,
446  la: Array<F32>, ra: Array<F32>) -> Model.Balanced(F32, depth, build(depth, splits, operation, plan, count, base, la, ra)):
447  Parallel.balanced(~Request, ~Walk, ~Array<F32>, ~Array<F32>, ~walk, ~split, ~cell_at, Request{operation, plan}, depth, splits, count, base, (la, ra))
448
449# The input handed back to the caller and an array copy for the build. A
450# tensor with empty storage has no cell to read.
451def take_owned(+shape: Shape.Shape, +depth: Nat, +length: Nat, -array: Array<F32>,
452  pair: Array<F32> & Array<F32>, -same: {pair == Array.clone(F32, array) : Array<F32> & Array<F32>},
453  -balanced: Model.Balanced(F32, depth, array), -valid: {Storage.bounds(depth, length) == True{} : Bool},
454  -exact: {length == Shape.logical_numel(Shape.dimensions(shape)) : Nat}) -> Core.Dynamic & Array<F32>:
455  (kept, copy) = pair
456  (Core.Packed{shape, Core.TTensor{Storage.Owned{depth, length, kept,
457    Certificates.clone_first(F32, depth, array, (kept, copy), same, balanced), valid}, exact}}, copy)
458
459def take(dynamic: Core.Dynamic) -> Core.Dynamic & Array<F32>:
460  Core.Packed{+shape, Core.TTensor{storage, exact}} = dynamic
461  match storage:
462    case Storage.Empty{}:
463      (Core.Packed{shape, Core.TTensor{Storage.Empty{}, exact}}, Array.new(F32, 0n, 0.0))
464    case Storage.Owned{+depth, +length, array, balanced, valid}:
465      take_owned(shape, depth, length, array, Array.clone(F32, array), {==}, balanced, valid, exact)
466
467def parallel_output(+operation: Operation, +plan: Shape.BroadcastPlan, +depth: Nat, +length: Nat,
468  -valid: {Storage.bounds(depth, length) == True{} : Bool},
469  -exact: {length == Shape.logical_numel(Shape.dimensions(Shape.output_shape(plan))) : Nat},
470  left: Core.Dynamic & Array<F32>, right: Core.Dynamic & Array<F32>) ->
471  (Core.Dynamic & Core.Dynamic) & Result<&2, &1, Core.Error, Core.Dynamic>:
472  (kept_left, la) = left
473  (kept_right, ra) = right
474  ((kept_left, kept_right), Done{Core.pack(Shape.output_shape(plan), Core.TTensor{Storage.Owned{depth, length,
475    build(depth, splits(depth), operation, plan, length, 0n, la, ra),
476    build_balanced(depth, splits(depth), operation, plan, length, 0n, la, ra),
477    valid}, exact})})
478
479def parallel_allocated(+operation: Operation, +plan: Shape.BroadcastPlan, left: Core.Dynamic, right: Core.Dynamic,
480  result: Result<&2, &1, Core.Error, Core.Tensor<Shape.output_shape(plan)>>) ->
481  (Core.Dynamic & Core.Dynamic) & Result<&2, &1, Core.Error, Core.Dynamic>:
482  match result:
483    case Fail{error}:
484      ((left, right), Fail{error})
485    case Done{Core.TTensor{storage, exact}}:
486      match storage:
487        case Storage.Empty{}:
488          ((left, right), Done{Core.pack(Shape.output_shape(plan), Core.TTensor{Storage.Empty{}, exact})})
489        case Storage.Owned{+depth, +length, array, balanced, valid}:
490          parallel_output(operation, plan, depth, length, valid, exact, take(left), take(right))
491
492def parallel_planned(operation: Operation, left: Core.Dynamic, right: Core.Dynamic,
493  result: Result<&2, &2, Shape.Error, Shape.BroadcastPlan>) ->
494  (Core.Dynamic & Core.Dynamic) & Result<&2, &1, Core.Error, Core.Dynamic>:
495  match result:
496    case Fail{error}:
497      ((left, right), Fail{Core.ShapeError{error}})
498    case Done{+plan}:
499      parallel_allocated(operation, plan, left, right, Core.new(Shape.output_shape(plan), 0.0))
500
501def binary(operation: Operation, left: Core.Dynamic, right: Core.Dynamic) ->
502  (Core.Dynamic & Core.Dynamic) & Result<&2, &1, Core.Error, Core.Dynamic>:
503  Core.Packed{+left_shape, left_tensor} = left
504  Core.Packed{+right_shape, right_tensor} = right
505  parallel_planned(operation, Core.pack(left_shape, left_tensor), Core.pack(right_shape, right_tensor),
506    Shape.broadcast_checked(left_shape, right_shape))
507
508def add(left: Core.Dynamic, right: Core.Dynamic) ->
509  (Core.Dynamic & Core.Dynamic) & Result<&2, &1, Core.Error, Core.Dynamic>:
510  binary(Add{}, left, right)
511
512def sub(left: Core.Dynamic, right: Core.Dynamic) ->
513  (Core.Dynamic & Core.Dynamic) & Result<&2, &1, Core.Error, Core.Dynamic>:
514  binary(Sub{}, left, right)
515
516def mul(left: Core.Dynamic, right: Core.Dynamic) ->
517  (Core.Dynamic & Core.Dynamic) & Result<&2, &1, Core.Error, Core.Dynamic>:
518  binary(Mul{}, left, right)
519
520def step_observe(-shape: Shape.Shape, step: Step<shape>, index: U32) -> Maybe<&2, F32>:
521  match step:
522    case Stop{left, right, error}:
523      None{}
524    case Continue{State{left, right, output}}:
525      Core.observe(shape, output, index)
526
527def expected_right(operation: Operation, left: F32, right: Maybe<&2, F32>) -> Maybe<&2, F32>:
528  match right:
529    case None{}:
530      None{}
531    case Some{value}:
532      Some{scalar(operation, left, value)}
533
534def expected_left(operation: Operation, right: Core.Dynamic, right_index: U32, left: Maybe<&2, F32>) -> Maybe<&2, F32>:
535  match left:
536    case None{}:
537      None{}
538    case Some{value}:
539      expected_right(operation, value, Core.checked_observe_dynamic(right, right_index))
540
541def expected_indices(operation: Operation, left: Core.Dynamic, right: Core.Dynamic,
542  result: Result<&2, &1, Core.Error, U32 & U32>) -> Maybe<&2, F32>:
543  match result:
544    case Fail{error}:
545      None{}
546    case Done{indices}:
547      (left_index, right_index) = indices
548      expected_left(operation, right, right_index, Core.checked_observe_dynamic(left, left_index))
549
550def expected(+operation: Operation, +plan: Shape.BroadcastPlan, left: Core.Dynamic, right: Core.Dynamic, index: U32) -> Maybe<&2, F32>:
551  expected_indices(operation, left, right, input_indices(plan, index))
552
553def PrefixMapping(+operation: Operation, +plan: Shape.BroadcastPlan,
554  -left: Core.Dynamic, -right: Core.Dynamic, +next: Nat, -output: Core.Dynamic) -> Type:
555  @+index: U32 -> @earlier: Shape.LT(U32.to_nat(index), next) ->
556    {Core.observe_dynamic(output, index) == expected(operation, plan, left, right, index) : Maybe<&2, F32>}
557
558def PrefixInvariant(+operation: Operation, +plan: Shape.BroadcastPlan,
559  -left: Core.Dynamic, -right: Core.Dynamic, +next: Nat, -output: Core.Dynamic) -> Type:
560  @+index: U32 -> @earlier: Shape.LT(U32.to_nat(index), next) ->
561    Core.ObservationChecked(output, index) &
562      {Core.observe_dynamic(output, index) == expected(operation, plan, left, right, index) : Maybe<&2, F32>}
563
564def SuccessfulMapping(+operation: Operation, +plan: Shape.BroadcastPlan, -left: Core.Dynamic, -right: Core.Dynamic,
565  pair: (Core.Dynamic & Core.Dynamic) & Result<&2, &1, Core.Error, Core.Dynamic>) -> Type:
566  (inputs, result) = pair
567  match result:
568    case Fail{error}:
569      Unit
570    case Done{output}:
571      PrefixMapping(operation, plan, left, right, Shape.logical_numel(Shape.dimensions(Shape.output_shape(plan))), output)
572
573def SameSource(-original: Core.Dynamic, -retained: Core.Dynamic) -> Type:
574  @+other: U32 -> {Core.checked_observe_dynamic(retained, other) == Core.checked_observe_dynamic(original, other) : Maybe<&2, F32>}
575
576def StepMeaning(+operation: Operation, +plan: Shape.BroadcastPlan, -index: U32,
577  original: State<Shape.output_shape(plan)>, result: Step<Shape.output_shape(plan)>) -> Type:
578  State{original_left, original_right, original_output} = original
579  match result:
580    case Stop{left, right, error}:
581      Unit
582    case Continue{State{left, right, output}}:
583      SameSource(original_left, left) &
584        (SameSource(original_right, right) &
585          (Core.ObservationChecked(Core.Packed{Shape.output_shape(plan), output}, index) &
586            ({Core.observe(Shape.output_shape(plan), output, index) == expected(operation, plan, original_left, original_right, index) : Maybe<&2, F32>} &
587              (@+other: U32 -> @checked: Core.ObservationChecked(Core.Packed{Shape.output_shape(plan), original_output}, other) ->
588                @different: {Nat.is_eq(U32.to_nat(index), U32.to_nat(other)) == False{} : Bool} ->
589                  Core.ObservationChecked(Core.Packed{Shape.output_shape(plan), output}, other) &
590                    {Core.observe(Shape.output_shape(plan), output, other) == Core.observe(Shape.output_shape(plan), original_output, other) : Maybe<&2, F32>}))))
591
592type StepCertificate<-operation: Operation, -plan: Shape.BroadcastPlan, -index: U32, -original: State<Shape.output_shape(plan)>> is Type:
593  StepCertificate{result: Step<Shape.output_shape(plan)>,
594    exact: {result == step(operation, plan, index, original) : Step<Shape.output_shape(plan)>},
595    -meaning: StepMeaning(operation, plan, index, original, result)}
596
597def NaturalMeaning(+operation: Operation, +plan: Shape.BroadcastPlan, -cursor: Nat,
598  original: State<Shape.output_shape(plan)>, result: NaturalStep<Shape.output_shape(plan), cursor>) -> Type:
599  match result:
600    case NaturalStop{left, right, error}:
601      Unit
602    case NaturalContinue{state, index, exact}:
603      StepMeaning(operation, plan, index, original, Continue{state})
604
605type ConvertedCertificate<-operation: Operation, -plan: Shape.BroadcastPlan, -cursor: Nat, -index: U32,
606  -decision: Shape.BoolCase<Nat.is_eq(U32.to_nat(index), cursor)>, -original: State<Shape.output_shape(plan)>> is Type:
607  ConvertedCertificate{result: NaturalStep<Shape.output_shape(plan), cursor>,
608    exact: {result == natural_step_decided(operation, plan, cursor, index, original, decision) : NaturalStep<Shape.output_shape(plan), cursor>},
609    -meaning: NaturalMeaning(operation, plan, cursor, original, result)}
610
611type NaturalCertificate<-operation: Operation, -plan: Shape.BroadcastPlan, -cursor: Nat, -original: State<Shape.output_shape(plan)>> is Type:
612  NaturalCertificate{result: NaturalStep<Shape.output_shape(plan), cursor>,
613    exact: {result == natural_step(operation, plan, cursor, original) : NaturalStep<Shape.output_shape(plan), cursor>},
614    -meaning: NaturalMeaning(operation, plan, cursor, original, result)}
615
616def StateInvariant(+operation: Operation, +plan: Shape.BroadcastPlan, -original_left: Core.Dynamic, -original_right: Core.Dynamic,
617  +cursor: Nat, state: State<Shape.output_shape(plan)>) -> Type:
618  State{left, right, output} = state
619  SameSource(original_left, left) &
620    (SameSource(original_right, right) & PrefixInvariant(operation, plan, original_left, original_right, cursor, Core.Packed{Shape.output_shape(plan), output}))
621
622type InvariantModel<-operation: Operation, -plan: Shape.BroadcastPlan, -original_left: Core.Dynamic, -original_right: Core.Dynamic,
623  -cursor: Nat, -state: State<Shape.output_shape(plan)>> is Data:
624  InvariantModel{-invariant: StateInvariant(operation, plan, original_left, original_right, cursor, state)}
625
626def LoopMeaning(+operation: Operation, +plan: Shape.BroadcastPlan, -original_left: Core.Dynamic, -original_right: Core.Dynamic,
627  +limit: Nat, pair: (Core.Dynamic & Core.Dynamic) & Result<&2, &1, Core.Error, Core.Dynamic>) -> Type:
628  (inputs, result) = pair
629  match result:
630    case Fail{error}:
631      Unit
632    case Done{output}:
633      PrefixInvariant(operation, plan, original_left, original_right, limit, output)
634
635type LoopCertificate<-fuel: Nat, -operation: Operation, -plan: Shape.BroadcastPlan, -cursor: Nat,
636  -original_left: Core.Dynamic, -original_right: Core.Dynamic, -original: State<Shape.output_shape(plan)>> is Type:
637  LoopCertificate{result: (Core.Dynamic & Core.Dynamic) & Result<&2, &1, Core.Error, Core.Dynamic>,
638    exact: {result == loop(fuel, operation, plan, cursor, original) : (Core.Dynamic & Core.Dynamic) & Result<&2, &1, Core.Error, Core.Dynamic>},
639    -meaning: LoopMeaning(operation, plan, original_left, original_right, Nat.add(cursor, fuel), result)}
640
641type MappingModel<-operation: Operation, -plan: Shape.BroadcastPlan, -original_left: Core.Dynamic, -original_right: Core.Dynamic,
642  -result: (Core.Dynamic & Core.Dynamic) & Result<&2, &1, Core.Error, Core.Dynamic>> is Data:
643  MappingModel{-mapping: SuccessfulMapping(operation, plan, original_left, original_right, result)}
644
645def plan_for(left: Core.Dynamic, right: Core.Dynamic) -> Result<&2, &2, Shape.Error, Shape.BroadcastPlan>:
646  Core.Packed{left_shape, left_tensor} = left
647  Core.Packed{right_shape, right_tensor} = right
648  Shape.broadcast_checked(left_shape, right_shape)
649
650def PlannedMeaning(+operation: Operation, -left: Core.Dynamic, -right: Core.Dynamic,
651  result: (Core.Dynamic & Core.Dynamic) & Result<&2, &1, Core.Error, Core.Dynamic>,
652  plan: Result<&2, &2, Shape.Error, Shape.BroadcastPlan>) -> Type:
653  match plan:
654    case Fail{error}:
655      Unit
656    case Done{+plan}:
657      SuccessfulMapping(operation, plan, left, right, result)
658
659def BinaryMeaning(+operation: Operation, left: Core.Dynamic, right: Core.Dynamic,
660  result: (Core.Dynamic & Core.Dynamic) & Result<&2, &1, Core.Error, Core.Dynamic>) -> Type:
661  PlannedMeaning(operation, left, right, result, plan_for(left, right))
662
663type BinaryModel<-operation: Operation, -left: Core.Dynamic, -right: Core.Dynamic,
664  -result: (Core.Dynamic & Core.Dynamic) & Result<&2, &1, Core.Error, Core.Dynamic>> is Data:
665  BinaryModel{-mapping: BinaryMeaning(operation, left, right, result)}