src/tensor/slice_PROOF.bend

bendtorch-core@0.1.0.0 · 965 lines · 66 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/nat_proofs.bend as Nats
4import bendtorch-core@0.1.0.0/src/shape/index_proofs.bend as Index
5import bendtorch-core@0.1.0.0/src/storage/arith.bend as Arith
6import bendtorch-core@0.1.0.0/src/storage/api.bend as Api
7import bendtorch-core@0.1.0.0/src/storage/model.bend as Model
8import bendtorch-core@0.1.0.0/src/storage/refinement.bend as Refinement
9import bendtorch-core@0.1.0.0/src/tensor/core.bend as Core
10import bendtorch-core@0.1.0.0/src/tensor/layout.bend as Layout
11import bendtorch-core@0.1.0.0/src/tensor/transforms.bend as Transforms
12import bendtorch-core@0.1.0.0/src/tensor/slice.bend as Slice
13import bendtorch-core@0.1.0.0/src/tensor/slice_LAWS.bend as Laws
14import bendtorch-core@0.1.0.0/src/tensor/contracts/build_facts.bend as Facts
15import bendtorch-core@0.1.0.0/src/tensor/elementwise.bend as Elementwise
16import bendtorch-core@0.1.0.0/src/tensor/elementwise_parallel_proof.bend as Parallel
17import bendtorch-core@0.1.0.0/src/tensor/parallel_proofs.bend as Generic
18import bendtorch-core@0.1.0.0/src/tensor/batched_matmul_PROOF.bend as BatchedProof
19import bendtorch-core@0.1.0.0/src/storage/array.bend as ArrayCore
20import bendtorch-core@0.1.0.0/src/tensor/elementwise_walk_proof.bend as Walk
21import bendtorch-core@0.1.0.0/src/autograd/reduce_fiber_proof.bend as Fiber
22
23# The clamp.
24
25def min_left(+a: Nat, +b: Nat, h: {Nat.is_le(a, b) == True{} : Bool}) -> {Nat.min(a, b) == a : Nat}:
26  match a b:
27    case 0n _:
28      {==}
29    case 1n+p 0n:
30      Empty.absurd({Nat.min(1n+p, 0n) == 1n+p : Nat}, Arith.clash(h))
31    case 1n+p 1n+q:
32      Equal.cong(Nat, Nat, x => 1n+x, Nat.min(p, q), p, min_left(p, q, h))
33
34def min_right(+a: Nat, +b: Nat, h: {Nat.is_le(b, a) == True{} : Bool}) -> {Nat.min(a, b) == b : Nat}:
35  match a b:
36    case 0n 0n:
37      {==}
38    case 0n 1n+q:
39      Empty.absurd({Nat.min(0n, 1n+q) == 1n+q : Nat}, Arith.clash(h))
40    case 1n+p 0n:
41      {==}
42    case 1n+p 1n+q:
43      Equal.cong(Nat, Nat, x => 1n+x, Nat.min(p, q), q, min_right(p, q, h))
44
45def min_le_right(+a: Nat, +b: Nat) -> {Nat.is_le(Nat.min(a, b), b) == True{} : Bool}:
46  match a b:
47    case 0n _:
48      Arith.le_zero(b)
49    case 1n+p 0n:
50      {==}
51    case 1n+p 1n+q:
52      min_le_right(p, q)
53
54# min(a, b) <= min(c, b) for a <= c.
55def min_mono(+a: Nat, +c: Nat, +b: Nat, h: {Nat.is_le(a, c) == True{} : Bool}) -> {Nat.is_le(Nat.min(a, b), Nat.min(c, b)) == True{} : Bool}:
56  match a c b:
57    case 0n _ _:
58      Arith.le_zero(Nat.min(c, b))
59    case 1n+p 0n _:
60      Empty.absurd({Nat.is_le(Nat.min(1n+p, b), Nat.min(0n, b)) == True{} : Bool}, Arith.clash(h))
61    case 1n+p 1n+q 0n:
62      {==}
63    case 1n+p 1n+q 1n+r:
64      min_mono(p, q, r, h)
65
66def max_right(+a: Nat, +b: Nat, h: {Nat.is_le(a, b) == True{} : Bool}) -> {Nat.max(a, b) == b : Nat}:
67  match a b:
68    case 0n _:
69      {==}
70    case 1n+p 0n:
71      Empty.absurd({Nat.max(1n+p, 0n) == 0n : Nat}, Arith.clash(h))
72    case 1n+p 1n+q:
73      Equal.cong(Nat, Nat, x => 1n+x, Nat.max(p, q), q, max_right(p, q, h))
74
75def max_left(+a: Nat, +b: Nat, h: {Nat.is_le(b, a) == True{} : Bool}) -> {Nat.max(a, b) == a : Nat}:
76  match a b:
77    case 0n 0n:
78      {==}
79    case 0n 1n+q:
80      Empty.absurd({Nat.max(0n, 1n+q) == 0n : Nat}, Arith.clash(h))
81    case 1n+p 0n:
82      {==}
83    case 1n+p 1n+q:
84      Equal.cong(Nat, Nat, x => 1n+x, Nat.max(p, q), p, max_left(p, q, h))
85
86def Laws.slice_start_clamped(extent, start, stop, past):
87  min_right(start, extent, past)
88
89def Laws.slice_start_kept(extent, start, stop, inside):
90  min_left(start, extent, inside)
91
92def Laws.slice_stop_clamped(extent, start, stop, past):
93  %Equal.sym(Nat, Nat.min(stop, extent), extent, min_right(stop, extent, past)) :
94    {Nat.max(Nat.min(start, extent), _) == extent : Nat}
95  max_right(Nat.min(start, extent), extent, min_le_right(start, extent))
96
97def Laws.slice_stop_kept(extent, start, stop, after, inside):
98  %Equal.sym(Nat, Nat.min(start, extent), start, min_left(start, extent, Facts.le_trans(start, stop, extent, after, inside))) :
99    {Nat.max(_, Nat.min(stop, extent)) == stop : Nat}
100  %Equal.sym(Nat, Nat.min(stop, extent), stop, min_left(stop, extent, inside)) :
101    {Nat.max(start, _) == stop : Nat}
102  max_right(start, stop, after)
103
104def Laws.slice_stop_empty(extent, start, stop, before):
105  max_left(Nat.min(start, extent), Nat.min(stop, extent), min_mono(stop, start, extent, before))
106
107# The length.
108
109def add_zero_right(+number: Nat) -> {Nat.add(number, 0n) == number : Nat}:
110  match number:
111    case 0n:
112      {==}
113    case 1n+p:
114      Equal.cong(Nat, Nat, x => 1n+x, Nat.add(p, 0n), p, add_zero_right(p))
115
116# Nat facts the closed form of the length needs.
117
118def nat_clash(+x: Nat, e: {1n+x == 0n : Nat}) -> Empty:
119  Arith.clash(Equal.cong(Nat, Bool, v => Nat.is_eq(v, 0n), 1n+x, 0n, e))
120
121# stop - low = 1 + rest gives stop = low + (1 + rest).
122def sub_back(+stop: Nat, +low: Nat, +rest: Nat, e: {Nat.sub(stop, low) == 1n+rest : Nat}) -> {Nat.add(low, 1n+rest) == stop : Nat}:
123  match stop low:
124    case 0n _:
125      Empty.absurd({Nat.add(low, 1n+rest) == 0n : Nat}, nat_clash(rest, Equal.sym(Nat, 0n, 1n+rest, e)))
126    case 1n+ap 0n:
127      Equal.sym(Nat, 1n+ap, 1n+rest, e)
128    case 1n+ap 1n+lp:
129      Equal.cong(Nat, Nat, x => 1n+x, Nat.add(lp, 1n+rest), ap, sub_back(ap, lp, rest, e))
130
131# stop - low = 0 gives stop <= low.
132def sub_zero_le(+stop: Nat, +low: Nat, e: {Nat.sub(stop, low) == 0n : Nat}) -> {Nat.is_le(stop, low) == True{} : Bool}:
133  match stop low:
134    case 0n _:
135      Arith.le_zero(low)
136    case 1n+ap 0n:
137      Empty.absurd({Nat.is_le(1n+ap, 0n) == True{} : Bool}, nat_clash(ap, e))
138    case 1n+ap 1n+lp:
139      sub_zero_le(ap, lp, e)
140
141def lt_succ_le(+k: Nat, +q: Nat, h: {Nat.is_lt(k, 1n+q) == True{} : Bool}) -> {Nat.is_le(k, q) == True{} : Bool}:
142  match k q:
143    case 0n _:
144      Arith.le_zero(q)
145    case 1n+p 0n:
146      Empty.absurd({Nat.is_le(1n+p, 0n) == True{} : Bool}, Arith.lt_zero(p, h))
147    case 1n+p 1n+r:
148      lt_succ_le(p, r, h)
149
150def lt_le_succ(+r: Nat, +s: Nat, h: {Nat.is_lt(r, s) == True{} : Bool}) -> {Nat.is_le(1n+r, s) == True{} : Bool}:
151  match r s:
152    case 0n 0n:
153      Empty.absurd({Nat.is_le(1n, 0n) == True{} : Bool}, Arith.clash(h))
154    case 0n 1n+q:
155      Arith.le_zero(q)
156    case 1n+p 0n:
157      Empty.absurd({Nat.is_le(2n+p, 0n) == True{} : Bool}, Arith.clash(h))
158    case 1n+p 1n+q:
159      lt_le_succ(p, q, h)
160
161# k <= rest / step gives k * step <= rest; at step 0 the product is 0.
162def scaled_within(+k: Nat, +rest: Nat, +s: Nat, +k_le: {Nat.is_le(k, Nat.div(rest, 1n+s)) == True{} : Bool},
163  given: Fiber.Decomposed<rest, 1n+s>) -> {Nat.is_le(Nat.mul(k, 1n+s), rest) == True{} : Bool}:
164  Fiber.Decomposed{below, sum} = given
165  Fiber.le_trans(Nat.mul(k, 1n+s), Nat.mul(Nat.div(rest, 1n+s), 1n+s), rest,
166    Fiber.mul_le_mono(k, Nat.div(rest, 1n+s), 1n+s, k_le),
167    %sum : {Nat.is_le(Nat.mul(Nat.div(rest, 1n+s), 1n+s), _) == True{} : Bool}
168    Api.le_add_right(Nat.mul(Nat.div(rest, 1n+s), 1n+s), Nat.mod(rest, 1n+s)))
169
170def scaled_le(+k: Nat, +rest: Nat, +step: Nat, below: {Nat.is_lt(k, 1n+Nat.div(rest, step)) == True{} : Bool}) ->
171  {Nat.is_le(Nat.mul(k, step), rest) == True{} : Bool}:
172  match step:
173    case 0n:
174      %Equal.sym(Nat, Nat.mul(k, 0n), 0n, Nats.mul_zero_right(k)) : {Nat.is_le(_, rest) == True{} : Bool}
175      Arith.le_zero(rest)
176    case 1n+s:
177      scaled_within(k, rest, s, lt_succ_le(k, Nat.div(rest, 1n+s), below), Fiber.decompose(rest, 1n+s, {==}))
178
179def below_span(+start: Nat, +stop: Nat, +step: Nat, +k: Nat, +span: Nat, e: {Nat.sub(stop, start) == span : Nat},
180  below: {Nat.is_lt(k, Slice.counted(span, step)) == True{} : Bool}) ->
181  {Nat.is_lt(Nat.add(start, Nat.mul(k, step)), stop) == True{} : Bool}:
182  match span:
183    case 0n:
184      Empty.absurd({Nat.is_lt(Nat.add(start, Nat.mul(k, step)), stop) == True{} : Bool}, Arith.lt_zero(k, below))
185    case 1n+rest:
186      %sub_back(stop, start, rest, e) : {Nat.is_lt(Nat.add(start, Nat.mul(k, step)), _) == True{} : Bool}
187      %Equal.sym(Nat, Nat.add(start, 1n+rest), 1n+Nat.add(start, rest), Arith.add_succ(start, rest)) :
188        {Nat.is_lt(Nat.add(start, Nat.mul(k, step)), _) == True{} : Bool}
189      Arith.lt_succ(Nat.add(start, Nat.mul(k, step)), Nat.add(start, rest),
190        Facts.le_add_cong(start, Nat.mul(k, step), rest, scaled_le(k, rest, step, below)))
191
192def Laws.slice_length_below(start, stop, step, k, below):
193  below_span(start, stop, step, k, Nat.sub(stop, start), {==}, below)
194
195# 1 + rest <= step + (rest / step) * step for a positive step: rest's
196# remainder is below the step.
197def covered(+rest: Nat, +s: Nat, given: Fiber.Decomposed<rest, 1n+s>) ->
198  {Nat.is_le(1n+rest, Nat.add(1n+s, Nat.mul(Nat.div(rest, 1n+s), 1n+s))) == True{} : Bool}:
199  Fiber.Decomposed{below, sum} = given
200  +q = Nat.mul(Nat.div(rest, 1n+s), 1n+s)
201  %sum : {Nat.is_le(1n+_, Nat.add(1n+s, q)) == True{} : Bool}
202  %Walk.add_comm(q, 1n+s) : {Nat.is_le(1n+Nat.add(q, Nat.mod(rest, 1n+s)), _) == True{} : Bool}
203  %Arith.add_succ(q, Nat.mod(rest, 1n+s)) : {Nat.is_le(_, Nat.add(q, 1n+s)) == True{} : Bool}
204  Facts.le_add_cong(q, 1n+Nat.mod(rest, 1n+s), 1n+s, lt_le_succ(Nat.mod(rest, 1n+s), 1n+s, below))
205
206def cover(+rest: Nat, +step: Nat, positive: {Nat.is_lt(0n, step) == True{} : Bool}) ->
207  {Nat.is_le(1n+rest, Nat.add(step, Nat.mul(Nat.div(rest, step), step))) == True{} : Bool}:
208  match step:
209    case 0n:
210      Empty.absurd({Nat.is_le(1n+rest, Nat.add(0n, Nat.mul(Nat.div(rest, 0n), 0n))) == True{} : Bool}, Arith.clash(positive))
211    case 1n+s:
212      covered(rest, s, Fiber.decompose(rest, 1n+s, {==}))
213
214def reaches_span(+start: Nat, +stop: Nat, +step: Nat, +span: Nat, e: {Nat.sub(stop, start) == span : Nat},
215  positive: {Nat.is_lt(0n, step) == True{} : Bool}) ->
216  {Nat.is_le(stop, Nat.add(start, Nat.mul(Slice.counted(span, step), step))) == True{} : Bool}:
217  match span:
218    case 0n:
219      %Equal.sym(Nat, Nat.add(start, 0n), start, add_zero_right(start)) : {Nat.is_le(stop, _) == True{} : Bool}
220      sub_zero_le(stop, start, e)
221    case 1n+rest:
222      %sub_back(stop, start, rest, e) : {Nat.is_le(_, Nat.add(start, Nat.mul(1n+Nat.div(rest, step), step))) == True{} : Bool}
223      Facts.le_add_cong(start, 1n+rest, Nat.add(step, Nat.mul(Nat.div(rest, step), step)), cover(rest, step, positive))
224
225def Laws.slice_length_reaches(start, stop, step, positive):
226  reaches_span(start, stop, step, Nat.sub(stop, start), {==}, positive)
227
228# The view, the output dimensions and select's dimensions.
229
230def view_around(+prefix: List<&2, U32>, +extent: U32, +suffix: List<&2, U32>, +outer: Nat) ->
231  {Slice.view_at(List.append(&2, U32, prefix, extent <> suffix), Shape.rank_dimensions(prefix), outer) ==
232    Slice.View{Nat.mul(outer, Shape.logical_numel(prefix)), U32.to_nat(extent), Shape.logical_numel(suffix)} : Slice.View}:
233  match prefix:
234    case Nil{}:
235      %Equal.sym(Nat, Nat.mul(outer, 1n), outer, Nats.mul_one_right(outer)) :
236        {Slice.View{outer, U32.to_nat(extent), Shape.logical_numel(suffix)} ==
237          Slice.View{_, U32.to_nat(extent), Shape.logical_numel(suffix)} : Slice.View}
238      {==}
239    case +head <> +tail:
240      %Nats.mul_assoc(outer, U32.to_nat(head), Shape.logical_numel(tail)) :
241        {Slice.view_at(List.append(&2, U32, tail, extent <> suffix), Shape.rank_dimensions(tail), Nat.mul(outer, U32.to_nat(head))) ==
242          Slice.View{_, U32.to_nat(extent), Shape.logical_numel(suffix)} : Slice.View}
243      view_around(tail, extent, suffix, Nat.mul(outer, U32.to_nat(head)))
244
245def Laws.slice_view_around(prefix, extent, suffix):
246  Equal.trans(Slice.View, Slice.view(List.append(&2, U32, prefix, extent <> suffix), Shape.rank_dimensions(prefix)),
247    Slice.View{Nat.mul(1n, Shape.logical_numel(prefix)), U32.to_nat(extent), Shape.logical_numel(suffix)},
248    Slice.View{Shape.logical_numel(prefix), U32.to_nat(extent), Shape.logical_numel(suffix)},
249    view_around(prefix, extent, suffix, 1n),
250    Equal.cong(Nat, Slice.View, x => Slice.View{x, U32.to_nat(extent), Shape.logical_numel(suffix)},
251      Nat.mul(1n, Shape.logical_numel(prefix)), Shape.logical_numel(prefix), add_zero_right(Shape.logical_numel(prefix))))
252
253def Laws.slice_resized_around(prefix, extent, suffix, other):
254  match prefix:
255    case Nil{}:
256      {==}
257    case +head <> +tail:
258      Equal.cong(List<&2, U32>, List<&2, U32>, rest => head <> rest,
259        Slice.resized(List.append(&2, U32, tail, extent <> suffix), Shape.rank_dimensions(tail), other),
260        List.append(&2, U32, tail, other <> suffix), Laws.slice_resized_around(tail, extent, suffix, other))
261
262def Laws.select_dropped_around(prefix, extent, suffix):
263  match prefix:
264    case Nil{}:
265      {==}
266    case +head <> +tail:
267      Equal.cong(List<&2, U32>, List<&2, U32>, rest => head <> rest,
268        Slice.dropped(List.append(&2, U32, tail, extent <> suffix), Shape.rank_dimensions(tail)),
269        List.append(&2, U32, tail, suffix), Laws.select_dropped_around(tail, extent, suffix))
270
271# The plan from the request.
272
273def pick_succ(+c: Bool, +p: Nat, +q: Nat) -> {Bool.pick(Nat, c, 1n+p, 1n+q) == 1n+Bool.pick(Nat, c, p, q) : Nat}:
274  match c:
275    case True{}:
276      {==}
277    case False{}:
278      {==}
279
280def min_pick(+a: Nat, +b: Nat) -> {Nat.min(a, b) == Bool.pick(Nat, Nat.is_le(a, b), a, b) : Nat}:
281  match a b:
282    case 0n 0n:
283      {==}
284    case 0n 1n+q:
285      {==}
286    case 1n+p 0n:
287      {==}
288    case 1n+p 1n+q:
289      Equal.trans(Nat, 1n+Nat.min(p, q), 1n+Bool.pick(Nat, Nat.is_le(p, q), p, q), Bool.pick(Nat, Nat.is_le(p, q), 1n+p, 1n+q),
290        Equal.cong(Nat, Nat, x => 1n+x, Nat.min(p, q), Bool.pick(Nat, Nat.is_le(p, q), p, q), min_pick(p, q)),
291        Equal.sym(Nat, Bool.pick(Nat, Nat.is_le(p, q), 1n+p, 1n+q), 1n+Bool.pick(Nat, Nat.is_le(p, q), p, q), pick_succ(Nat.is_le(p, q), p, q)))
292
293def max_pick(+a: Nat, +b: Nat) -> {Nat.max(a, b) == Bool.pick(Nat, Nat.is_le(a, b), b, a) : Nat}:
294  match a b:
295    case 0n 0n:
296      {==}
297    case 0n 1n+q:
298      {==}
299    case 1n+p 0n:
300      {==}
301    case 1n+p 1n+q:
302      Equal.trans(Nat, 1n+Nat.max(p, q), 1n+Bool.pick(Nat, Nat.is_le(p, q), q, p), Bool.pick(Nat, Nat.is_le(p, q), 1n+q, 1n+p),
303        Equal.cong(Nat, Nat, x => 1n+x, Nat.max(p, q), Bool.pick(Nat, Nat.is_le(p, q), q, p), max_pick(p, q)),
304        Equal.sym(Nat, Bool.pick(Nat, Nat.is_le(p, q), 1n+q, 1n+p), 1n+Bool.pick(Nat, Nat.is_le(p, q), q, p), pick_succ(Nat.is_le(p, q), q, p)))
305
306def view_extent(+given: Slice.View) -> Nat:
307  Slice.View{outer, extent, inner} = given
308  extent
309
310def view_pair(+given: Slice.View) -> Nat & Nat:
311  Slice.View{outer, extent, inner} = given
312  (outer, inner)
313
314# The clamp of bounds is PyTorch's, written with min and max.
315def clamp_axis(+n: Nat, +s: Nat, +t: Nat, +p: Nat) ->
316  {(n, Nat.min(s, n), p, Slice.length(Slice.Bounds{Nat.min(s, n), Nat.max(Nat.min(s, n), Nat.min(t, n))}, p)) ==
317    Laws.axis_spec(n, s, t, p) : Nat & Nat & Nat & Nat}:
318  +cs = Bool.pick(Nat, Nat.is_le(s, n), s, n)
319  +ct = Bool.pick(Nat, Nat.is_le(t, n), t, n)
320  %Equal.sym(Nat, Nat.min(s, n), cs, min_pick(s, n)) :
321    {(n, _, p, Slice.length(Slice.Bounds{_, Nat.max(_, Nat.min(t, n))}, p)) == Laws.axis_spec(n, s, t, p) : Nat & Nat & Nat & Nat}
322  %Equal.sym(Nat, Nat.min(t, n), ct, min_pick(t, n)) :
323    {(n, cs, p, Slice.length(Slice.Bounds{cs, Nat.max(cs, _)}, p)) == Laws.axis_spec(n, s, t, p) : Nat & Nat & Nat & Nat}
324  %Equal.sym(Nat, Nat.max(cs, ct), Bool.pick(Nat, Nat.is_le(cs, ct), ct, cs), max_pick(cs, ct)) :
325    {(n, cs, p, Slice.length(Slice.Bounds{cs, _}, p)) == Laws.axis_spec(n, s, t, p) : Nat & Nat & Nat & Nat}
326  {==}
327
328def viewed_axis(+given: Slice.View, +s: Nat, +t: Nat, +p: Nat) ->
329  {Laws.axis_part(Slice.viewed(given, s, t, p)) == Laws.axis_spec(view_extent(given), s, t, p) : Nat & Nat & Nat & Nat}:
330  match given:
331    case Slice.View{+o, +n, +i}:
332      clamp_axis(n, s, t, p)
333
334def viewed_view(+given: Slice.View, +s: Nat, +t: Nat, +p: Nat) ->
335  {Laws.view_part(Slice.viewed(given, s, t, p)) == view_pair(given) : Nat & Nat}:
336  match given:
337    case Slice.View{+o, +n, +i}:
338      {==}
339
340def axis_extent_around(+prefix: List<&2, U32>, +extent: U32, +suffix: List<&2, U32>) ->
341  {Slice.axis_extent(List.append(&2, U32, prefix, extent <> suffix), Shape.rank_dimensions(prefix)) == U32.to_nat(extent) : Nat}:
342  match prefix:
343    case Nil{}:
344      {==}
345    case +head <> +tail:
346      axis_extent_around(tail, extent, suffix)
347
348def sized_extent(+prefix: List<&2, U32>, +extent: U32, +suffix: List<&2, U32>, +empty: Bool) ->
349  {view_extent(Slice.sized(List.append(&2, U32, prefix, extent <> suffix), Shape.rank_dimensions(prefix), empty)) == U32.to_nat(extent) : Nat}:
350  match empty:
351    case True{}:
352      axis_extent_around(prefix, extent, suffix)
353    case False{}:
354      Equal.cong(Slice.View, Nat, v => view_extent(v),
355        Slice.view(List.append(&2, U32, prefix, extent <> suffix), Shape.rank_dimensions(prefix)),
356        Slice.View{Shape.logical_numel(prefix), U32.to_nat(extent), Shape.logical_numel(suffix)},
357        Laws.slice_view_around(prefix, extent, suffix))
358
359def Laws.slice_plan_axis(prefix, extent, suffix, dim, at, start, stop, step):
360  +dims = List.append(&2, U32, prefix, extent <> suffix)
361  %Equal.sym(Nat, U32.to_nat(dim), Shape.rank_dimensions(prefix), at) :
362    {Laws.axis_part(Slice.viewed(Slice.sized(dims, _, Slice.holds_zero(dims)), U32.to_nat(start), U32.to_nat(stop), U32.to_nat(step))) ==
363      Laws.axis_spec(U32.to_nat(extent), U32.to_nat(start), U32.to_nat(stop), U32.to_nat(step)) : Nat & Nat & Nat & Nat}
364  Equal.trans(Nat & Nat & Nat & Nat,
365    Laws.axis_part(Slice.viewed(Slice.sized(dims, Shape.rank_dimensions(prefix), Slice.holds_zero(dims)),
366      U32.to_nat(start), U32.to_nat(stop), U32.to_nat(step))),
367    Laws.axis_spec(view_extent(Slice.sized(dims, Shape.rank_dimensions(prefix), Slice.holds_zero(dims))),
368      U32.to_nat(start), U32.to_nat(stop), U32.to_nat(step)),
369    Laws.axis_spec(U32.to_nat(extent), U32.to_nat(start), U32.to_nat(stop), U32.to_nat(step)),
370    viewed_axis(Slice.sized(dims, Shape.rank_dimensions(prefix), Slice.holds_zero(dims)), U32.to_nat(start), U32.to_nat(stop), U32.to_nat(step)),
371    Equal.cong(Nat, Nat & Nat & Nat & Nat, n => Laws.axis_spec(n, U32.to_nat(start), U32.to_nat(stop), U32.to_nat(step)),
372      view_extent(Slice.sized(dims, Shape.rank_dimensions(prefix), Slice.holds_zero(dims))), U32.to_nat(extent),
373      sized_extent(prefix, extent, suffix, Slice.holds_zero(dims))))
374
375# Some dimension 0 makes the element count 0.
376def zero_head(+extent: U32, +rest: List<&2, U32>, h: {Slice.holds_zero(extent <> rest) == True{} : Bool},
377  rec: @h: {Slice.holds_zero(rest) == True{} : Bool} -> {Shape.logical_numel(rest) == 0n : Nat},
378  decision: Shape.BoolCase<Nat.is_eq(U32.to_nat(extent), 0n)>) -> {Shape.logical_numel(extent <> rest) == 0n : Nat}:
379  match decision:
380    case Shape.BTrue{equal}:
381      %Equal.sym(Nat, U32.to_nat(extent), 0n, Arith.eq_true(U32.to_nat(extent), 0n, Equal.sym(Bool, True{}, Nat.is_eq(U32.to_nat(extent), 0n), equal))) :
382        {Nat.mul(_, Shape.logical_numel(rest)) == 0n : Nat}
383      {==}
384    case Shape.BFalse{equal}:
385      %Equal.sym(Nat, Shape.logical_numel(rest), 0n,
386        rec(Equal.trans(Bool, Slice.holds_zero(rest), Bool.or(Nat.is_eq(U32.to_nat(extent), 0n), Slice.holds_zero(rest)), True{},
387          Equal.cong(Bool, Bool, x => Bool.or(x, Slice.holds_zero(rest)), False{}, Nat.is_eq(U32.to_nat(extent), 0n), equal), h))) :
388        {Nat.mul(U32.to_nat(extent), _) == 0n : Nat}
389      Nats.mul_zero_right(U32.to_nat(extent))
390
391def zero_numel(+dimensions: List<&2, U32>, h: {Slice.holds_zero(dimensions) == True{} : Bool}) ->
392  {Shape.logical_numel(dimensions) == 0n : Nat}:
393  match dimensions:
394    case Nil{}:
395      Empty.absurd({Shape.logical_numel(Nil{}) == 0n : Nat}, Arith.clash(h))
396    case +extent <> +rest:
397      zero_head(extent, rest, h, hr => zero_numel(rest, hr), Shape.bool_case(Nat.is_eq(U32.to_nat(extent), 0n)))
398
399# A product a * b of 0 with a not 0 has b = 0.
400def add_zero_left_part(+b: Nat, +x: Nat, e: {Nat.add(b, x) == 0n : Nat}) -> {b == 0n : Nat}:
401  match b:
402    case 0n:
403      {==}
404    case 1n+q:
405      Empty.absurd({1n+q == 0n : Nat}, nat_clash(Nat.add(q, x), e))
406
407def mul_zero_factor(+a: Nat, +b: Nat, e: {Nat.mul(a, b) == 0n : Nat}, nonzero: {False{} == Nat.is_eq(a, 0n) : Bool}) -> {b == 0n : Nat}:
408  match a:
409    case 0n:
410      Empty.absurd({b == 0n : Nat}, Arith.clash(nonzero))
411    case 1n+p:
412      add_zero_left_part(b, Nat.mul(p, b), e)
413
414# An element count of 0 has some dimension 0.
415def zero_found(+extent: U32, +rest: List<&2, U32>, e: {Shape.logical_numel(extent <> rest) == 0n : Nat},
416  rec: @e: {Shape.logical_numel(rest) == 0n : Nat} -> {Slice.holds_zero(rest) == True{} : Bool},
417  decision: Shape.BoolCase<Nat.is_eq(U32.to_nat(extent), 0n)>) -> {Slice.holds_zero(extent <> rest) == True{} : Bool}:
418  match decision:
419    case Shape.BTrue{equal}:
420      %equal : {Bool.or(_, Slice.holds_zero(rest)) == True{} : Bool}
421      {==}
422    case Shape.BFalse{+equal}:
423      %equal : {Bool.or(_, Slice.holds_zero(rest)) == True{} : Bool}
424      rec(mul_zero_factor(U32.to_nat(extent), Shape.logical_numel(rest), e, equal))
425
426def holds_true(+dimensions: List<&2, U32>, e: {Shape.logical_numel(dimensions) == 0n : Nat}) -> {Slice.holds_zero(dimensions) == True{} : Bool}:
427  match dimensions:
428    case Nil{}:
429      Empty.absurd({Slice.holds_zero(Nil{}) == True{} : Bool}, nat_clash(0n, e))
430    case +extent <> +rest:
431      zero_found(extent, rest, e, er => holds_true(rest, er), Shape.bool_case(Nat.is_eq(U32.to_nat(extent), 0n)))
432
433def holds_false(+dimensions: List<&2, U32>, +nonempty: {Nat.is_lt(0n, Shape.logical_numel(dimensions)) == True{} : Bool},
434  decision: Shape.BoolCase<Slice.holds_zero(dimensions)>) -> {Slice.holds_zero(dimensions) == False{} : Bool}:
435  match decision:
436    case Shape.BFalse{equal}:
437      Equal.sym(Bool, False{}, Slice.holds_zero(dimensions), equal)
438    case Shape.BTrue{equal}:
439      Empty.absurd({Slice.holds_zero(dimensions) == False{} : Bool},
440        Arith.clash(Equal.trans(Bool, False{}, Nat.is_lt(0n, Shape.logical_numel(dimensions)), True{},
441          Equal.cong(Nat, Bool, v => Nat.is_lt(0n, v), 0n, Shape.logical_numel(dimensions),
442            Equal.sym(Nat, Shape.logical_numel(dimensions), 0n, zero_numel(dimensions, Equal.sym(Bool, True{}, Slice.holds_zero(dimensions), equal)))),
443          nonempty)))
444
445def Laws.slice_plan_view(prefix, extent, suffix, dim, at, start, stop, step, nonempty):
446  +dims = List.append(&2, U32, prefix, extent <> suffix)
447  %Equal.sym(Nat, U32.to_nat(dim), Shape.rank_dimensions(prefix), at) :
448    {Laws.view_part(Slice.viewed(Slice.sized(dims, _, Slice.holds_zero(dims)), U32.to_nat(start), U32.to_nat(stop), U32.to_nat(step))) ==
449      (Shape.logical_numel(prefix), Shape.logical_numel(suffix)) : Nat & Nat}
450  %Equal.sym(Bool, Slice.holds_zero(dims), False{}, holds_false(dims, nonempty, Shape.bool_case(Slice.holds_zero(dims)))) :
451    {Laws.view_part(Slice.viewed(Slice.sized(dims, Shape.rank_dimensions(prefix), _), U32.to_nat(start), U32.to_nat(stop), U32.to_nat(step))) ==
452      (Shape.logical_numel(prefix), Shape.logical_numel(suffix)) : Nat & Nat}
453  %Equal.sym(Slice.View, Slice.view(dims, Shape.rank_dimensions(prefix)),
454    Slice.View{Shape.logical_numel(prefix), U32.to_nat(extent), Shape.logical_numel(suffix)}, Laws.slice_view_around(prefix, extent, suffix)) :
455    {Laws.view_part(Slice.viewed(_, U32.to_nat(start), U32.to_nat(stop), U32.to_nat(step))) ==
456      (Shape.logical_numel(prefix), Shape.logical_numel(suffix)) : Nat & Nat}
457  {==}
458
459def Laws.slice_plan_empty_view(dimensions, dim, start, stop, step, empty):
460  %Equal.sym(Bool, Slice.holds_zero(dimensions), True{}, holds_true(dimensions, empty)) :
461    {Laws.view_part(Slice.viewed(Slice.sized(dimensions, U32.to_nat(dim), _), U32.to_nat(start), U32.to_nat(stop), U32.to_nat(step))) ==
462      (0n, 0n) : Nat & Nat}
463  {==}
464
465# The checks and the plan.
466
467def Laws.slice_axis_refused(input, dim, start, stop, step, outside):
468  %Equal.sym(Bool, Nat.is_lt(U32.to_nat(dim), Shape.rank_dimensions(Shape.dimensions(input))), False{}, outside) :
469    {Slice.checked(Shape.dimensions(input), dim, start, stop, step, Slice.refused(_, Nat.is_lt(0n, U32.to_nat(step)))) ==
470      Fail{Slice.LayoutError{Layout.AxisOutOfBounds{}}} : Result<&2, &2, Slice.Error, Slice.Planned>}
471  {==}
472
473def Laws.slice_step_refused(input, dim, start, stop, step, inside, zero):
474  %Equal.sym(Bool, Nat.is_lt(U32.to_nat(dim), Shape.rank_dimensions(Shape.dimensions(input))), True{}, inside) :
475    {Slice.checked(Shape.dimensions(input), dim, start, stop, step, Slice.refused(_, Nat.is_lt(0n, U32.to_nat(step)))) ==
476      Fail{Slice.StepNotPositive{}} : Result<&2, &2, Slice.Error, Slice.Planned>}
477  %Equal.sym(Nat, U32.to_nat(step), 0n, zero) :
478    {Slice.checked(Shape.dimensions(input), dim, start, stop, step, Slice.refused(True{}, Nat.is_lt(0n, _))) ==
479      Fail{Slice.StepNotPositive{}} : Result<&2, &2, Slice.Error, Slice.Planned>}
480  {==}
481
482def planned_shaped(+dimensions: List<&2, U32>, +dim: U32, +start: U32, +stop: U32, +step: U32,
483  built: Shape.ConstructResult<Slice.resized(dimensions, U32.to_nat(dim), U32.from_nat(Slice.plan_length(Slice.plan_of(dimensions, dim, start, stop, step))))>) ->
484  Laws.PlannedAs(dimensions, dim, start, stop, step,
485    Slice.shaped(Slice.plan_of(dimensions, dim, start, stop, step),
486      Slice.resized(dimensions, U32.to_nat(dim), U32.from_nat(Slice.plan_length(Slice.plan_of(dimensions, dim, start, stop, step)))), built)):
487  match built:
488    case Shape.SError{error}:
489      Unit{}
490    case Shape.SBuilt{+shape, preserved}:
491      ({==}, preserved)
492
493def planned_checked(+dimensions: List<&2, U32>, +dim: U32, +start: U32, +stop: U32, +step: U32, refused: Maybe<&2, Slice.Error>) ->
494  Laws.PlannedAs(dimensions, dim, start, stop, step, Slice.checked(dimensions, dim, start, stop, step, refused)):
495  match refused:
496    case Some{error}:
497      Unit{}
498    case None{}:
499      planned_shaped(dimensions, dim, start, stop, step,
500        Shape.construct(Slice.resized(dimensions, U32.to_nat(dim), U32.from_nat(Slice.plan_length(Slice.plan_of(dimensions, dim, start, stop, step))))))
501
502def Laws.slice_planned(input, dim, start, stop, step):
503  planned_checked(Shape.dimensions(input), dim, start, stop, step, Slice.refusal(Shape.dimensions(input), dim, step))
504
505def Laws.slice_refused(input, dim, start, stop, step, error, refused):
506  Core.Packed{+shape, tensor} = input
507  %Equal.sym(Result<&2, &2, Slice.Error, Slice.Planned>, Slice.plan(shape, dim, start, stop, step), Fail{error}, refused) :
508    {Slice.sliced(shape, tensor, _) == (Core.Packed{shape, tensor}, Fail{error}) : Core.Dynamic & Result<&2, &1, Slice.Error, Core.Dynamic>}
509  {==}
510
511# narrow and select.
512
513def Laws.narrow_out_of_range(input, dim, start, length, inside, past):
514  Core.Packed{+shape, tensor} = input
515  %Equal.sym(Bool, Nat.is_lt(U32.to_nat(dim), Shape.rank_dimensions(Shape.dimensions(shape))), True{}, inside) :
516    {Slice.narrow_checked(shape, tensor, dim, start, length, Slice.refused(_, Nat.is_lt(0n, U32.to_nat(1)))) ==
517      (Core.Packed{shape, tensor}, Fail{Slice.OutOfRange{}}) : Core.Dynamic & Result<&2, &1, Slice.Error, Core.Dynamic>}
518  %Equal.sym(Bool, Nat.is_le(Nat.add(U32.to_nat(start), U32.to_nat(length)), Slice.extent_at(Shape.dimensions(shape), dim)), False{}, past) :
519    {Slice.narrow_ranged(shape, tensor, dim, start, length, _) ==
520      (Core.Packed{shape, tensor}, Fail{Slice.OutOfRange{}}) : Core.Dynamic & Result<&2, &1, Slice.Error, Core.Dynamic>}
521  {==}
522
523def Laws.narrow_slices(input, dim, start, length, inside, fits):
524  Core.Packed{+shape, tensor} = input
525  %Equal.sym(Bool, Nat.is_lt(U32.to_nat(dim), Shape.rank_dimensions(Shape.dimensions(shape))), True{}, inside) :
526    {Slice.narrow_checked(shape, tensor, dim, start, length, Slice.refused(_, Nat.is_lt(0n, U32.to_nat(1)))) ==
527      Slice.slice(Core.Packed{shape, tensor}, dim, start, U32.add(start, length), 1) : Core.Dynamic & Result<&2, &1, Slice.Error, Core.Dynamic>}
528  %Equal.sym(Bool, Nat.is_le(Nat.add(U32.to_nat(start), U32.to_nat(length)), Slice.extent_at(Shape.dimensions(shape), dim)), True{}, fits) :
529    {Slice.narrow_ranged(shape, tensor, dim, start, length, _) ==
530      Slice.slice(Core.Packed{shape, tensor}, dim, start, U32.add(start, length), 1) : Core.Dynamic & Result<&2, &1, Slice.Error, Core.Dynamic>}
531  {==}
532
533def Laws.narrow_axis_refused(input, dim, start, length, outside):
534  Core.Packed{+shape, tensor} = input
535  %Equal.sym(Bool, Nat.is_lt(U32.to_nat(dim), Shape.rank_dimensions(Shape.dimensions(shape))), False{}, outside) :
536    {Slice.narrow_checked(shape, tensor, dim, start, length, Slice.refused(_, Nat.is_lt(0n, U32.to_nat(1)))) ==
537      (Core.Packed{shape, tensor}, Fail{Slice.LayoutError{Layout.AxisOutOfBounds{}}}) : Core.Dynamic & Result<&2, &1, Slice.Error, Core.Dynamic>}
538  {==}
539
540def transformed_as(result: Result<&2, &1, Transforms.Error, Core.Dynamic>) ->
541  {Slice.transformed(result) == Laws.reshape_as_slice(result) : Result<&2, &1, Slice.Error, Core.Dynamic>}:
542  match result:
543    case Fail{Transforms.LayoutError{error}}:
544      {==}
545    case Fail{Transforms.TensorError{error}}:
546      {==}
547    case Done{output}:
548      {==}
549
550def reshaped_as(+dimensions: List<&2, U32>, pair: Core.Dynamic & Result<&2, &1, Slice.Error, Core.Dynamic>) ->
551  {Slice.reshaped(dimensions, pair) == Laws.reshaped(dimensions, pair) : Core.Dynamic & Result<&2, &1, Slice.Error, Core.Dynamic>}:
552  match pair:
553    case Tuple{input, Fail{error}}:
554      {==}
555    case Tuple{input, Done{output}}:
556      Equal.cong(Result<&2, &1, Slice.Error, Core.Dynamic>, Core.Dynamic & Result<&2, &1, Slice.Error, Core.Dynamic>,
557        result => (input, result), Slice.transformed(Transforms.reshape_dynamic(output, dimensions)),
558        Laws.reshape_as_slice(Transforms.reshape_dynamic(output, dimensions)), transformed_as(Transforms.reshape_dynamic(output, dimensions)))
559
560def Laws.select_out_of_range(input, dim, index, inside, past):
561  Core.Packed{+shape, tensor} = input
562  %Equal.sym(Bool, Nat.is_lt(U32.to_nat(dim), Shape.rank_dimensions(Shape.dimensions(shape))), True{}, inside) :
563    {Slice.select_checked(shape, tensor, dim, index, Slice.refused(_, Nat.is_lt(0n, U32.to_nat(1)))) ==
564      (Core.Packed{shape, tensor}, Fail{Slice.OutOfRange{}}) : Core.Dynamic & Result<&2, &1, Slice.Error, Core.Dynamic>}
565  %Equal.sym(Bool, Nat.is_lt(U32.to_nat(index), Slice.extent_at(Shape.dimensions(shape), dim)), False{}, past) :
566    {Slice.select_ranged(shape, tensor, dim, index, _) ==
567      (Core.Packed{shape, tensor}, Fail{Slice.OutOfRange{}}) : Core.Dynamic & Result<&2, &1, Slice.Error, Core.Dynamic>}
568  {==}
569
570def Laws.select_slices(input, dim, index, inside, below):
571  Core.Packed{+shape, tensor} = input
572  %Equal.sym(Bool, Nat.is_lt(U32.to_nat(dim), Shape.rank_dimensions(Shape.dimensions(shape))), True{}, inside) :
573    {Slice.select_checked(shape, tensor, dim, index, Slice.refused(_, Nat.is_lt(0n, U32.to_nat(1)))) ==
574      Laws.reshaped(Slice.dropped(Shape.dimensions(shape), U32.to_nat(dim)), Slice.slice(Core.Packed{shape, tensor}, dim, index, U32.add(index, 1), 1)) :
575      Core.Dynamic & Result<&2, &1, Slice.Error, Core.Dynamic>}
576  %Equal.sym(Bool, Nat.is_lt(U32.to_nat(index), Slice.extent_at(Shape.dimensions(shape), dim)), True{}, below) :
577    {Slice.select_ranged(shape, tensor, dim, index, _) ==
578      Laws.reshaped(Slice.dropped(Shape.dimensions(shape), U32.to_nat(dim)), Slice.slice(Core.Packed{shape, tensor}, dim, index, U32.add(index, 1), 1)) :
579      Core.Dynamic & Result<&2, &1, Slice.Error, Core.Dynamic>}
580  reshaped_as(Slice.dropped(Shape.dimensions(shape), U32.to_nat(dim)), Slice.slice(Core.Packed{shape, tensor}, dim, index, U32.add(index, 1), 1))
581
582def Laws.select_axis_refused(input, dim, index, outside):
583  Core.Packed{+shape, tensor} = input
584  %Equal.sym(Bool, Nat.is_lt(U32.to_nat(dim), Shape.rank_dimensions(Shape.dimensions(shape))), False{}, outside) :
585    {Slice.select_checked(shape, tensor, dim, index, Slice.refused(_, Nat.is_lt(0n, U32.to_nat(1)))) ==
586      (Core.Packed{shape, tensor}, Fail{Slice.LayoutError{Layout.AxisOutOfBounds{}}}) : Core.Dynamic & Result<&2, &1, Slice.Error, Core.Dynamic>}
587  {==}
588
589# The cells of the builds.
590
591def take_value(-b: Array<F32>, pair: Array<F32> & F32) ->
592  {Pair.fst(F32, Array<F32> & Array<F32>, Slice.take(b, pair)) == Refinement.second(F32, pair) : F32}:
593  (a, x) = pair
594  {==}
595
596def take_state(-b: Array<F32>, pair: Array<F32> & F32) ->
597  {Pair.snd(F32, Array<F32> & Array<F32>, Slice.take(b, pair)) == (Refinement.first(F32, pair), b) : Array<F32> & Array<F32>}:
598  (a, x) = pair
599  {==}
600
601# A read of A through `take` gives the word A holds there ...
602def read_word(+da: Nat, -A: Array<F32>, -b: Array<F32>, +index: Nat, -valid: Model.Balanced(F32, da, A)) ->
603  Parallel.Fact<{Some{Pair.fst(F32, Array<F32> & Array<F32>, Slice.take(b, Array.get(F32, A, U32.from_nat(index))))} ==
604    Model.observe(F32, da, A, Refinement.route(da, U32.from_nat(index))) : Maybe<&2, F32>}>:
605  Parallel.Fact{
606    %Equal.sym(F32, Pair.fst(F32, Array<F32> & Array<F32>, Slice.take(b, Array.get(F32, A, U32.from_nat(index)))),
607      Refinement.second(F32, Array.get(F32, A, U32.from_nat(index))), take_value(b, Array.get(F32, A, U32.from_nat(index)))) :
608      {Some{_} == Model.observe(F32, da, A, Refinement.route(da, U32.from_nat(index))) : Maybe<&2, F32>}
609    {\{Parallel.Fact: q => q} :
610      Parallel.Fact<{Some{Refinement.second(F32, Array.get(F32, A, U32.from_nat(index)))} ==
611        Model.observe(F32, da, A, Refinement.route(da, U32.from_nat(index))) : Maybe<&2, F32>}> ->
612      {Some{Refinement.second(F32, Array.get(F32, A, U32.from_nat(index)))} ==
613        Model.observe(F32, da, A, Refinement.route(da, U32.from_nat(index))) : Maybe<&2, F32>}}(
614      Parallel.read_value(da, A, A, U32.from_nat(index), Array.get(F32, A, U32.from_nat(index)), {==}, {==}, valid))}
615
616# ... and hands both arrays back.
617def read_keeps(+da: Nat, -A: Array<F32>, -b: Array<F32>, +index: Nat, -valid: Model.Balanced(F32, da, A)) ->
618  Parallel.Fact<{Pair.snd(F32, Array<F32> & Array<F32>, Slice.take(b, Array.get(F32, A, U32.from_nat(index)))) == (A, b) :
619    Array<F32> & Array<F32>}>:
620  Parallel.Fact{Equal.trans(Array<F32> & Array<F32>, Pair.snd(F32, Array<F32> & Array<F32>, Slice.take(b, Array.get(F32, A, U32.from_nat(index)))),
621    (Refinement.first(F32, Array.get(F32, A, U32.from_nat(index))), b), (A, b),
622    take_state(b, Array.get(F32, A, U32.from_nat(index))),
623    Equal.cong(Array<F32>, Array<F32> & Array<F32>, x => (x, b), Refinement.first(F32, Array.get(F32, A, U32.from_nat(index))), A,
624      {\{Parallel.Fact: q => q} :
625        Parallel.Fact<{Refinement.first(F32, Array.get(F32, A, U32.from_nat(index))) == A : Array<F32>}> ->
626        {Refinement.first(F32, Array.get(F32, A, U32.from_nat(index))) == A : Array<F32>}}(
627        Parallel.read_keeps(da, A, A, U32.from_nat(index), Array.get(F32, A, U32.from_nat(index)), {==}, {==}, valid))))}
628
629def empty() -> Array<F32>:
630  Array.new(F32, 0n, 0.0)
631
632def cell_keeps(+plan: Slice.Plan, +at: Nat, +da: Nat, -A: Array<F32>, inside: Bool, -valid: Model.Balanced(F32, da, A)) ->
633  Parallel.Fact<{Pair.snd(F32, Array<F32> & Array<F32>, Slice.cell(plan, at, (A, empty()), inside)) == (A, empty()) : Array<F32> & Array<F32>}>:
634  match inside:
635    case False{}:
636      Parallel.Fact{{==}}
637    case True{}:
638      read_keeps(da, A, empty(), Slice.source(plan, at), valid)
639
640def scattered_keeps(+da: Nat, -A: Array<F32>, at: Maybe<&2, Nat>, -valid: Model.Balanced(F32, da, A)) ->
641  Parallel.Fact<{Pair.snd(F32, Array<F32> & Array<F32>, Slice.scattered((A, empty()), at)) == (A, empty()) : Array<F32> & Array<F32>}>:
642  match at:
643    case None{}:
644      Parallel.Fact{{==}}
645    case Some{+index}:
646      read_keeps(da, A, empty(), index, valid)
647
648def scatter_keeps(+plan: Slice.Plan, +at: Nat, +da: Nat, -A: Array<F32>, inside: Bool, -valid: Model.Balanced(F32, da, A)) ->
649  Parallel.Fact<{Pair.snd(F32, Array<F32> & Array<F32>, Slice.scatter_cell(plan, at, (A, empty()), inside)) == (A, empty()) : Array<F32> & Array<F32>}>:
650  match inside:
651    case False{}:
652      Parallel.Fact{{==}}
653    case True{}:
654      scattered_keeps(da, A, Slice.target(plan, at), valid)
655
656def SplitFacts(-A: Array<F32>) -> Type:
657  {Pair.fst(Array<F32> & Array<F32>, Array<F32> & Array<F32>, Elementwise.split((A, empty()))) == (A, empty()) : Array<F32> & Array<F32>} &
658    {Pair.snd(Array<F32> & Array<F32>, Array<F32> & Array<F32>, Elementwise.split((A, empty()))) == (A, empty()) : Array<F32> & Array<F32>}
659
660def split_facts(+da: Nat, -A: Array<F32>, -valid: Model.Balanced(F32, da, A)) -> Parallel.Fact<SplitFacts(A)>:
661  Parallel.Fact{{\{Parallel.Fact: q => q} : Parallel.Fact<SplitFacts(A)> -> SplitFacts(A)}(
662    Parallel.split_equal(da, 0n, A, empty(), valid, ArrayCore.new_balanced(F32, 0n, 0.0)))}
663
664# Every slot of the forward build holds what its cell computes from (A, []).
665# The Fact holds the cells eta-expanded: `--verdict` rejects the unexpanded
666# function in this position although `--check-only` accepts it.
667def build_cells(+depth: Nat, +splits: Nat, +plan: Slice.Plan, +count: Nat, +base: Nat, +da: Nat, -A: Array<F32>, a: Array<F32>,
668  -same: {a == A : Array<F32>}, -valid: Model.Balanced(F32, da, A)) ->
669  Parallel.Fact<Generic.Cells(~Slice.Plan, ~Array<F32>, ~Array<F32>, ~Slice.cell, plan, (A, empty()), depth, count, base,
670    Slice.build(depth, splits, plan, count, base, a))>:
671  Parallel.Fact{path =>
672    {\{Generic.Fact: q => q} :
673      Generic.Fact<Generic.Cells(~Slice.Plan, ~Array<F32>, ~Array<F32>, ~Slice.cell, plan, (A, empty()), depth, count, base,
674        Slice.build(depth, splits, plan, count, base, a))> ->
675      Generic.Cells(~Slice.Plan, ~Array<F32>, ~Array<F32>, ~Slice.cell, plan, (A, empty()), depth, count, base,
676        Slice.build(depth, splits, plan, count, base, a))}(
677      Generic.build_cells(~Slice.Plan, ~Slice.Plan, ~Array<F32>, ~Array<F32>, ~Slice.same, ~Elementwise.split, ~Slice.cell,
678        plan, (A, empty()), depth, splits, count, base, (a, empty()), Parallel.pair_equal(a, empty(), A, empty(), same, {==}),
679        at => inside => {\{Parallel.Fact: q => q} :
680          Parallel.Fact<{Pair.snd(F32, Array<F32> & Array<F32>, Slice.cell(plan, at, (A, empty()), inside)) == (A, empty()) : Array<F32> & Array<F32>}> ->
681          {Pair.snd(F32, Array<F32> & Array<F32>, Slice.cell(plan, at, (A, empty()), inside)) == (A, empty()) : Array<F32> & Array<F32>}}(
682          cell_keeps(plan, at, da, A, inside, valid)),
683        Pair.fst({Pair.fst(Array<F32> & Array<F32>, Array<F32> & Array<F32>, Elementwise.split((A, empty()))) == (A, empty()) : Array<F32> & Array<F32>},
684          {Pair.snd(Array<F32> & Array<F32>, Array<F32> & Array<F32>, Elementwise.split((A, empty()))) == (A, empty()) : Array<F32> & Array<F32>},
685          {\{Parallel.Fact: q => q} : Parallel.Fact<SplitFacts(A)> -> SplitFacts(A)}(split_facts(da, A, valid))),
686        Pair.snd({Pair.fst(Array<F32> & Array<F32>, Array<F32> & Array<F32>, Elementwise.split((A, empty()))) == (A, empty()) : Array<F32> & Array<F32>},
687          {Pair.snd(Array<F32> & Array<F32>, Array<F32> & Array<F32>, Elementwise.split((A, empty()))) == (A, empty()) : Array<F32> & Array<F32>},
688          {\{Parallel.Fact: q => q} : Parallel.Fact<SplitFacts(A)> -> SplitFacts(A)}(split_facts(da, A, valid)))))(path)}
689
690def scatter_cells(+depth: Nat, +splits: Nat, +plan: Slice.Plan, +count: Nat, +base: Nat, +da: Nat, -A: Array<F32>, a: Array<F32>,
691  -same: {a == A : Array<F32>}, -valid: Model.Balanced(F32, da, A)) ->
692  Parallel.Fact<Generic.Cells(~Slice.Plan, ~Array<F32>, ~Array<F32>, ~Slice.scatter_cell, plan, (A, empty()), depth, count, base,
693    Slice.scatter_build(depth, splits, plan, count, base, a))>:
694  Parallel.Fact{path =>
695    {\{Generic.Fact: q => q} :
696      Generic.Fact<Generic.Cells(~Slice.Plan, ~Array<F32>, ~Array<F32>, ~Slice.scatter_cell, plan, (A, empty()), depth, count, base,
697        Slice.scatter_build(depth, splits, plan, count, base, a))> ->
698      Generic.Cells(~Slice.Plan, ~Array<F32>, ~Array<F32>, ~Slice.scatter_cell, plan, (A, empty()), depth, count, base,
699        Slice.scatter_build(depth, splits, plan, count, base, a))}(
700      Generic.build_cells(~Slice.Plan, ~Slice.Plan, ~Array<F32>, ~Array<F32>, ~Slice.same, ~Elementwise.split, ~Slice.scatter_cell,
701        plan, (A, empty()), depth, splits, count, base, (a, empty()), Parallel.pair_equal(a, empty(), A, empty(), same, {==}),
702        at => inside => {\{Parallel.Fact: q => q} :
703          Parallel.Fact<{Pair.snd(F32, Array<F32> & Array<F32>, Slice.scatter_cell(plan, at, (A, empty()), inside)) == (A, empty()) : Array<F32> & Array<F32>}> ->
704          {Pair.snd(F32, Array<F32> & Array<F32>, Slice.scatter_cell(plan, at, (A, empty()), inside)) == (A, empty()) : Array<F32> & Array<F32>}}(
705          scatter_keeps(plan, at, da, A, inside, valid)),
706        Pair.fst({Pair.fst(Array<F32> & Array<F32>, Array<F32> & Array<F32>, Elementwise.split((A, empty()))) == (A, empty()) : Array<F32> & Array<F32>},
707          {Pair.snd(Array<F32> & Array<F32>, Array<F32> & Array<F32>, Elementwise.split((A, empty()))) == (A, empty()) : Array<F32> & Array<F32>},
708          {\{Parallel.Fact: q => q} : Parallel.Fact<SplitFacts(A)> -> SplitFacts(A)}(split_facts(da, A, valid))),
709        Pair.snd({Pair.fst(Array<F32> & Array<F32>, Array<F32> & Array<F32>, Elementwise.split((A, empty()))) == (A, empty()) : Array<F32> & Array<F32>},
710          {Pair.snd(Array<F32> & Array<F32>, Array<F32> & Array<F32>, Elementwise.split((A, empty()))) == (A, empty()) : Array<F32> & Array<F32>},
711          {\{Parallel.Fact: q => q} : Parallel.Fact<SplitFacts(A)> -> SplitFacts(A)}(split_facts(da, A, valid)))))(path)}
712
713# The arithmetic of the positions.
714
715def sub_zero(+number: Nat) -> {Nat.sub(number, 0n) == number : Nat}:
716  match number:
717    case 0n:
718      {==}
719    case 1n+p:
720      {==}
721
722def sub_cancel(+a: Nat, +b: Nat) -> {Nat.sub(Nat.add(a, b), a) == b : Nat}:
723  match a:
724    case 0n:
725      sub_zero(b)
726    case 1n+p:
727      sub_cancel(p, b)
728
729# (k * step) / step = k for a positive step.
730def div_step(+k: Nat, +step: Nat, +positive: {Nat.is_lt(0n, step) == True{} : Bool}) -> {Nat.div(Nat.mul(k, step), step) == k : Nat}:
731  Equal.trans(Nat, Nat.div(Nat.mul(k, step), step), Nat.div(Nat.add(Nat.mul(k, step), 0n), step), k,
732    Equal.cong(Nat, Nat, x => Nat.div(x, step), Nat.mul(k, step), Nat.add(Nat.mul(k, step), 0n),
733      Equal.sym(Nat, Nat.add(Nat.mul(k, step), 0n), Nat.mul(k, step), add_zero_right(Nat.mul(k, step)))),
734    BatchedProof.div_exact(k, step, 0n, positive))
735
736# The forward cell (o, k, r) reads input cell (o, start + k * step, r).
737def source_at(+outer: Nat, +extent: Nat, +inner: Nat, +start: Nat, +step: Nat, +length: Nat, +o: Nat, +k: Nat, +r: Nat,
738  +k_below: {Nat.is_lt(k, length) == True{} : Bool}, +r_below: {Nat.is_lt(r, inner) == True{} : Bool}) ->
739  {Slice.source(Slice.Plan{outer, extent, inner, start, step, length}, Laws.position(length, inner, o, k, r)) ==
740    Laws.position(extent, inner, o, Nat.add(start, Nat.mul(k, step)), r) : Nat}:
741  +row = Nat.add(Nat.mul(o, length), k)
742  +pos = Laws.position(length, inner, o, k, r)
743  %Equal.sym(Nat, Nat.div(pos, inner), row, BatchedProof.div_exact(row, inner, r, r_below)) :
744    {Nat.add(Nat.mul(Nat.add(Nat.mul(Nat.div(_, length), extent), Nat.add(start, Nat.mul(Nat.mod(_, length), step))), inner), Nat.mod(pos, inner)) ==
745      Laws.position(extent, inner, o, Nat.add(start, Nat.mul(k, step)), r) : Nat}
746  %Equal.sym(Nat, Nat.mod(pos, inner), r, BatchedProof.mod_exact(row, inner, r, r_below)) :
747    {Nat.add(Nat.mul(Nat.add(Nat.mul(Nat.div(row, length), extent), Nat.add(start, Nat.mul(Nat.mod(row, length), step))), inner), _) ==
748      Laws.position(extent, inner, o, Nat.add(start, Nat.mul(k, step)), r) : Nat}
749  %Equal.sym(Nat, Nat.div(row, length), o, BatchedProof.div_exact(o, length, k, k_below)) :
750    {Nat.add(Nat.mul(Nat.add(Nat.mul(_, extent), Nat.add(start, Nat.mul(Nat.mod(row, length), step))), inner), r) ==
751      Laws.position(extent, inner, o, Nat.add(start, Nat.mul(k, step)), r) : Nat}
752  %Equal.sym(Nat, Nat.mod(row, length), k, BatchedProof.mod_exact(o, length, k, k_below)) :
753    {Nat.add(Nat.mul(Nat.add(Nat.mul(o, extent), Nat.add(start, Nat.mul(_, step))), inner), r) ==
754      Laws.position(extent, inner, o, Nat.add(start, Nat.mul(k, step)), r) : Nat}
755  {==}
756
757def Laws.slice_build_cells(depth, splits, outer, extent, inner, start, step, length, count, base, da, A, a, same, valid, path, o, k, r,
758  k_below, r_below, at, inside):
759  Laws.Holds{
760    %Equal.trans(Nat, Slice.source(Slice.Plan{outer, extent, inner, start, step, length}, Nat.add(Model.offset(depth, path), base)),
761      Slice.source(Slice.Plan{outer, extent, inner, start, step, length}, Laws.position(length, inner, o, k, r)),
762      Laws.position(extent, inner, o, Nat.add(start, Nat.mul(k, step)), r),
763      Equal.cong(Nat, Nat, x => Slice.source(Slice.Plan{outer, extent, inner, start, step, length}, x),
764        Nat.add(Model.offset(depth, path), base), Laws.position(length, inner, o, k, r), at),
765      source_at(outer, extent, inner, start, step, length, o, k, r, k_below, r_below)) :
766      {Model.observe(F32, depth, Slice.build(depth, splits, Slice.Plan{outer, extent, inner, start, step, length}, count, base, a), path) ==
767        Model.observe(F32, da, A, Refinement.route(da, U32.from_nat(_))) : Maybe<&2, F32>}
768    Equal.trans(Maybe<&2, F32>,
769      Model.observe(F32, depth, Slice.build(depth, splits, Slice.Plan{outer, extent, inner, start, step, length}, count, base, a), path),
770      Some{Pair.fst(F32, Array<F32> & Array<F32>, Slice.cell(Slice.Plan{outer, extent, inner, start, step, length},
771        Nat.add(Model.offset(depth, path), base), (A, empty()), True{}))},
772      Model.observe(F32, da, A, Refinement.route(da, U32.from_nat(Slice.source(Slice.Plan{outer, extent, inner, start, step, length},
773        Nat.add(Model.offset(depth, path), base))))),
774      %inside :
775        {Model.observe(F32, depth, Slice.build(depth, splits, Slice.Plan{outer, extent, inner, start, step, length}, count, base, a), path) ==
776          Some{Pair.fst(F32, Array<F32> & Array<F32>, Slice.cell(Slice.Plan{outer, extent, inner, start, step, length},
777            Nat.add(Model.offset(depth, path), base), (A, empty()), _))} : Maybe<&2, F32>}
778      {\{Parallel.Fact: q => q} :
779        Parallel.Fact<Generic.Cells(~Slice.Plan, ~Array<F32>, ~Array<F32>, ~Slice.cell, Slice.Plan{outer, extent, inner, start, step, length},
780          (A, empty()), depth, count, base, Slice.build(depth, splits, Slice.Plan{outer, extent, inner, start, step, length}, count, base, a))> ->
781        Generic.Cells(~Slice.Plan, ~Array<F32>, ~Array<F32>, ~Slice.cell, Slice.Plan{outer, extent, inner, start, step, length},
782          (A, empty()), depth, count, base, Slice.build(depth, splits, Slice.Plan{outer, extent, inner, start, step, length}, count, base, a))}(
783        build_cells(depth, splits, Slice.Plan{outer, extent, inner, start, step, length}, count, base, da, A, a, same, valid))(path),
784      {\{Parallel.Fact: q => q} :
785        Parallel.Fact<{Some{Pair.fst(F32, Array<F32> & Array<F32>, Slice.take(empty(), Array.get(F32, A, U32.from_nat(Slice.source(Slice.Plan{outer, extent, inner, start, step, length},
786          Nat.add(Model.offset(depth, path), base))))))} ==
787          Model.observe(F32, da, A, Refinement.route(da, U32.from_nat(Slice.source(Slice.Plan{outer, extent, inner, start, step, length},
788            Nat.add(Model.offset(depth, path), base))))) : Maybe<&2, F32>}> ->
789        {Some{Pair.fst(F32, Array<F32> & Array<F32>, Slice.take(empty(), Array.get(F32, A, U32.from_nat(Slice.source(Slice.Plan{outer, extent, inner, start, step, length},
790          Nat.add(Model.offset(depth, path), base))))))} ==
791          Model.observe(F32, da, A, Refinement.route(da, U32.from_nat(Slice.source(Slice.Plan{outer, extent, inner, start, step, length},
792            Nat.add(Model.offset(depth, path), base))))) : Maybe<&2, F32>}}(
793        read_word(da, A, empty(), Slice.source(Slice.Plan{outer, extent, inner, start, step, length}, Nat.add(Model.offset(depth, path), base)), valid)))}
794
795# The backward target of input cell (o, i, r).
796
797def picked_k(+outer: Nat, +extent: Nat, +inner: Nat, +start: Nat, +step: Nat, +length: Nat, +row: Nat, +within: Nat, +i: Nat, +k: Nat) ->
798  Maybe<&2, Nat>:
799  Slice.picked(Slice.Plan{outer, extent, inner, start, step, length}, row, within, i, k, Slice.hits(start, step, length, i, k))
800
801def picked_at(+outer: Nat, +extent: Nat, +inner: Nat, +start: Nat, +step: Nat, +length: Nat, +row: Nat, +within: Nat, +i: Nat) ->
802  Maybe<&2, Nat>:
803  picked_k(outer, extent, inner, start, step, length, row, within, i, Nat.div(Nat.sub(i, start), step))
804
805# `Slice.target` from the row and the remainder of the position.
806def target_of(+outer: Nat, +extent: Nat, +inner: Nat, +start: Nat, +step: Nat, +length: Nat, +row: Nat, +within: Nat) -> Maybe<&2, Nat>:
807  picked_at(outer, extent, inner, start, step, length, row, within, Nat.mod(row, extent))
808
809# Position (o * extent + i) * inner + r splits into row o * extent + i and
810# remainder r, and the row into o and i.
811def target_split(+outer: Nat, +extent: Nat, +inner: Nat, +start: Nat, +step: Nat, +length: Nat, +o: Nat, +i: Nat, +r: Nat,
812  +i_below: {Nat.is_lt(i, extent) == True{} : Bool}, +r_below: {Nat.is_lt(r, inner) == True{} : Bool}) ->
813  {Slice.target(Slice.Plan{outer, extent, inner, start, step, length}, Laws.position(extent, inner, o, i, r)) ==
814    picked_at(outer, extent, inner, start, step, length, Nat.add(Nat.mul(o, extent), i), r, i) : Maybe<&2, Nat>}:
815  +row = Nat.add(Nat.mul(o, extent), i)
816  +pos = Laws.position(extent, inner, o, i, r)
817  %Equal.sym(Nat, Nat.div(pos, inner), row, BatchedProof.div_exact(row, inner, r, r_below)) :
818    {target_of(outer, extent, inner, start, step, length, _, Nat.mod(pos, inner)) ==
819      picked_at(outer, extent, inner, start, step, length, row, r, i) : Maybe<&2, Nat>}
820  %Equal.sym(Nat, Nat.mod(pos, inner), r, BatchedProof.mod_exact(row, inner, r, r_below)) :
821    {target_of(outer, extent, inner, start, step, length, row, _) == picked_at(outer, extent, inner, start, step, length, row, r, i) : Maybe<&2, Nat>}
822  %Equal.sym(Nat, Nat.mod(row, extent), i, BatchedProof.mod_exact(o, extent, i, i_below)) :
823    {picked_at(outer, extent, inner, start, step, length, row, r, _) == picked_at(outer, extent, inner, start, step, length, row, r, i) : Maybe<&2, Nat>}
824  {==}
825
826# Input cell (o, i, r) that no k below the length reaches reads nothing.
827def missed_k(+outer: Nat, +extent: Nat, +inner: Nat, +start: Nat, +step: Nat, +length: Nat, +row: Nat, +r: Nat, +i: Nat, +k: Nat,
828  missed: @+k: Nat -> @below: {Nat.is_lt(k, length) == True{} : Bool} -> {Nat.is_eq(Nat.add(start, Nat.mul(k, step)), i) == False{} : Bool},
829  decision: Shape.BoolCase<Slice.hits(start, step, length, i, k)>) ->
830  {picked_k(outer, extent, inner, start, step, length, row, r, i, k) == None{} : Maybe<&2, Nat>}:
831  match decision:
832    case Shape.BFalse{equal}:
833      %equal : {Slice.picked(Slice.Plan{outer, extent, inner, start, step, length}, row, r, i, k, _) == None{} : Maybe<&2, Nat>}
834      {==}
835    case Shape.BTrue{equal}:
836      +rest = Arith.and_right(Nat.is_le(start, i), Bool.and(Nat.is_lt(k, length), Nat.is_eq(Nat.add(start, Nat.mul(k, step)), i)),
837        Equal.sym(Bool, True{}, Slice.hits(start, step, length, i, k), equal))
838      Empty.absurd({picked_k(outer, extent, inner, start, step, length, row, r, i, k) == None{} : Maybe<&2, Nat>},
839        Arith.clash(Equal.trans(Bool, False{}, Nat.is_eq(Nat.add(start, Nat.mul(k, step)), i), True{},
840          Equal.sym(Bool, Nat.is_eq(Nat.add(start, Nat.mul(k, step)), i), False{},
841            missed(k, Arith.and_left(Nat.is_lt(k, length), Nat.is_eq(Nat.add(start, Nat.mul(k, step)), i), rest))),
842          Arith.and_right(Nat.is_lt(k, length), Nat.is_eq(Nat.add(start, Nat.mul(k, step)), i), rest))))
843
844def target_miss(+outer: Nat, +extent: Nat, +inner: Nat, +start: Nat, +step: Nat, +length: Nat, +o: Nat, +i: Nat, +r: Nat,
845  +i_below: {Nat.is_lt(i, extent) == True{} : Bool}, +r_below: {Nat.is_lt(r, inner) == True{} : Bool},
846  missed: @+k: Nat -> @below: {Nat.is_lt(k, length) == True{} : Bool} -> {Nat.is_eq(Nat.add(start, Nat.mul(k, step)), i) == False{} : Bool}) ->
847  {Slice.target(Slice.Plan{outer, extent, inner, start, step, length}, Laws.position(extent, inner, o, i, r)) == None{} : Maybe<&2, Nat>}:
848  Equal.trans(Maybe<&2, Nat>, Slice.target(Slice.Plan{outer, extent, inner, start, step, length}, Laws.position(extent, inner, o, i, r)),
849    picked_at(outer, extent, inner, start, step, length, Nat.add(Nat.mul(o, extent), i), r, i), None{},
850    target_split(outer, extent, inner, start, step, length, o, i, r, i_below, r_below),
851    missed_k(outer, extent, inner, start, step, length, Nat.add(Nat.mul(o, extent), i), r, i, Nat.div(Nat.sub(i, start), step), missed,
852      Shape.bool_case(Slice.hits(start, step, length, i, Nat.div(Nat.sub(i, start), step)))))
853
854def ScatterCells(+depth: Nat, +splits: Nat, +plan: Slice.Plan, +count: Nat, +base: Nat, -G: Array<F32>, -g: Array<F32>) -> Type:
855  Generic.Cells(~Slice.Plan, ~Array<F32>, ~Array<F32>, ~Slice.scatter_cell, plan, (G, empty()), depth, count, base,
856    Slice.scatter_build(depth, splits, plan, count, base, g))
857
858# The slot of input cell (o, i, r), below the count, holds what the gradient
859# target of the cell names.
860def scatter_slot(+depth: Nat, +splits: Nat, +plan: Slice.Plan, +count: Nat, +base: Nat, +dg: Nat, -G: Array<F32>, g: Array<F32>,
861  -same: {g == G : Array<F32>}, -valid: Model.Balanced(F32, dg, G), +path: Model.Path(depth),
862  -inside: {Nat.is_lt(Nat.add(Model.offset(depth, path), base), count) == True{} : Bool}) ->
863  Parallel.Fact<{Model.observe(F32, depth, Slice.scatter_build(depth, splits, plan, count, base, g), path) ==
864    Some{Pair.fst(F32, Array<F32> & Array<F32>, Slice.scattered((G, empty()), Slice.target(plan, Nat.add(Model.offset(depth, path), base))))} :
865    Maybe<&2, F32>}>:
866  Parallel.Fact{
867    %inside :
868      {Model.observe(F32, depth, Slice.scatter_build(depth, splits, plan, count, base, g), path) ==
869        Some{Pair.fst(F32, Array<F32> & Array<F32>, Slice.scatter_cell(plan, Nat.add(Model.offset(depth, path), base), (G, empty()), _))} :
870        Maybe<&2, F32>}
871    {\{Parallel.Fact: q => q} : Parallel.Fact<ScatterCells(depth, splits, plan, count, base, G, g)> -> ScatterCells(depth, splits, plan, count, base, G, g)}(
872      scatter_cells(depth, splits, plan, count, base, dg, G, g, same, valid))(path)}
873
874def Laws.slice_scatter_miss(depth, splits, outer, extent, inner, start, step, length, count, base, dg, G, g, same, valid, path, o, i, r,
875  i_below, r_below, missed, at, inside):
876  Laws.Holds{
877    Equal.trans(Maybe<&2, F32>,
878      Model.observe(F32, depth, Slice.scatter_build(depth, splits, Slice.Plan{outer, extent, inner, start, step, length}, count, base, g), path),
879      Some{Pair.fst(F32, Array<F32> & Array<F32>, Slice.scattered((G, empty()),
880        Slice.target(Slice.Plan{outer, extent, inner, start, step, length}, Nat.add(Model.offset(depth, path), base))))},
881      Some{0.0},
882      {\{Parallel.Fact: q => q} :
883        Parallel.Fact<{Model.observe(F32, depth, Slice.scatter_build(depth, splits, Slice.Plan{outer, extent, inner, start, step, length}, count, base, g), path) ==
884          Some{Pair.fst(F32, Array<F32> & Array<F32>, Slice.scattered((G, empty()),
885            Slice.target(Slice.Plan{outer, extent, inner, start, step, length}, Nat.add(Model.offset(depth, path), base))))} : Maybe<&2, F32>}> ->
886        {Model.observe(F32, depth, Slice.scatter_build(depth, splits, Slice.Plan{outer, extent, inner, start, step, length}, count, base, g), path) ==
887          Some{Pair.fst(F32, Array<F32> & Array<F32>, Slice.scattered((G, empty()),
888            Slice.target(Slice.Plan{outer, extent, inner, start, step, length}, Nat.add(Model.offset(depth, path), base))))} : Maybe<&2, F32>}}(
889        scatter_slot(depth, splits, Slice.Plan{outer, extent, inner, start, step, length}, count, base, dg, G, g, same, valid, path, inside)),
890      %Equal.sym(Maybe<&2, Nat>, Slice.target(Slice.Plan{outer, extent, inner, start, step, length}, Nat.add(Model.offset(depth, path), base)), None{},
891        Equal.trans(Maybe<&2, Nat>, Slice.target(Slice.Plan{outer, extent, inner, start, step, length}, Nat.add(Model.offset(depth, path), base)),
892          Slice.target(Slice.Plan{outer, extent, inner, start, step, length}, Laws.position(extent, inner, o, i, r)), None{},
893          Equal.cong(Nat, Maybe<&2, Nat>, x => Slice.target(Slice.Plan{outer, extent, inner, start, step, length}, x),
894            Nat.add(Model.offset(depth, path), base), Laws.position(extent, inner, o, i, r), at),
895          target_miss(outer, extent, inner, start, step, length, o, i, r, i_below, r_below, missed))) :
896        {Some{Pair.fst(F32, Array<F32> & Array<F32>, Slice.scattered((G, empty()), _))} == Some{0.0} : Maybe<&2, F32>}
897      {==})}
898
899# start <= start + k * step, k < length and start + k * step is itself.
900def hits_true(+start: Nat, +step: Nat, +length: Nat, +k: Nat, +k_below: {Nat.is_lt(k, length) == True{} : Bool}) ->
901  {Slice.hits(start, step, length, Nat.add(start, Nat.mul(k, step)), k) == True{} : Bool}:
902  %Equal.sym(Bool, Nat.is_le(start, Nat.add(start, Nat.mul(k, step))), True{}, Api.le_add_right(start, Nat.mul(k, step))) :
903    {Bool.and(_, Bool.and(Nat.is_lt(k, length), Nat.is_eq(Nat.add(start, Nat.mul(k, step)), Nat.add(start, Nat.mul(k, step))))) == True{} : Bool}
904  %Equal.sym(Bool, Nat.is_lt(k, length), True{}, k_below) :
905    {Bool.and(True{}, Bool.and(_, Nat.is_eq(Nat.add(start, Nat.mul(k, step)), Nat.add(start, Nat.mul(k, step))))) == True{} : Bool}
906  %Equal.sym(Bool, Nat.is_eq(Nat.add(start, Nat.mul(k, step)), Nat.add(start, Nat.mul(k, step))), True{},
907    Arith.eq_refl(Nat.add(start, Nat.mul(k, step)))) :
908    {Bool.and(True{}, Bool.and(True{}, _)) == True{} : Bool}
909  {==}
910
911# Input cell (o, start + k * step, r) reads gradient cell (o, k, r).
912def target_hit(+outer: Nat, +extent: Nat, +inner: Nat, +start: Nat, +step: Nat, +length: Nat, +o: Nat, +k: Nat, +r: Nat,
913  +positive: {Nat.is_lt(0n, step) == True{} : Bool}, +k_below: {Nat.is_lt(k, length) == True{} : Bool},
914  +i_below: {Nat.is_lt(Nat.add(start, Nat.mul(k, step)), extent) == True{} : Bool}, +r_below: {Nat.is_lt(r, inner) == True{} : Bool}) ->
915  {Slice.target(Slice.Plan{outer, extent, inner, start, step, length}, Laws.position(extent, inner, o, Nat.add(start, Nat.mul(k, step)), r)) ==
916    Some{Laws.position(length, inner, o, k, r)} : Maybe<&2, Nat>}:
917  +i = Nat.add(start, Nat.mul(k, step))
918  +row = Nat.add(Nat.mul(o, extent), i)
919  Equal.trans(Maybe<&2, Nat>, Slice.target(Slice.Plan{outer, extent, inner, start, step, length}, Laws.position(extent, inner, o, i, r)),
920    picked_at(outer, extent, inner, start, step, length, row, r, i), Some{Laws.position(length, inner, o, k, r)},
921    target_split(outer, extent, inner, start, step, length, o, i, r, i_below, r_below),
922    %Equal.sym(Nat, Nat.div(Nat.sub(i, start), step), k,
923      Equal.trans(Nat, Nat.div(Nat.sub(i, start), step), Nat.div(Nat.mul(k, step), step), k,
924        Equal.cong(Nat, Nat, x => Nat.div(x, step), Nat.sub(i, start), Nat.mul(k, step), sub_cancel(start, Nat.mul(k, step))),
925        div_step(k, step, positive))) :
926      {picked_k(outer, extent, inner, start, step, length, row, r, i, _) == Some{Laws.position(length, inner, o, k, r)} : Maybe<&2, Nat>}
927    %Equal.sym(Bool, Slice.hits(start, step, length, i, k), True{}, hits_true(start, step, length, k, k_below)) :
928      {Slice.picked(Slice.Plan{outer, extent, inner, start, step, length}, row, r, i, k, _) == Some{Laws.position(length, inner, o, k, r)} :
929        Maybe<&2, Nat>}
930    %Equal.sym(Nat, Nat.div(row, extent), o, BatchedProof.div_exact(o, extent, i, i_below)) :
931      {Some{Nat.add(Nat.mul(Nat.add(Nat.mul(_, length), k), inner), r)} == Some{Laws.position(length, inner, o, k, r)} : Maybe<&2, Nat>}
932    {==})
933
934def Laws.slice_scatter_hit(depth, splits, outer, extent, inner, start, step, length, count, base, dg, G, g, same, valid, path, o, k, r,
935  positive, k_below, i_below, r_below, at, inside):
936  Laws.Holds{
937    Equal.trans(Maybe<&2, F32>,
938      Model.observe(F32, depth, Slice.scatter_build(depth, splits, Slice.Plan{outer, extent, inner, start, step, length}, count, base, g), path),
939      Some{Pair.fst(F32, Array<F32> & Array<F32>, Slice.scattered((G, empty()),
940        Slice.target(Slice.Plan{outer, extent, inner, start, step, length}, Nat.add(Model.offset(depth, path), base))))},
941      Model.observe(F32, dg, G, Refinement.route(dg, U32.from_nat(Laws.position(length, inner, o, k, r)))),
942      {\{Parallel.Fact: q => q} :
943        Parallel.Fact<{Model.observe(F32, depth, Slice.scatter_build(depth, splits, Slice.Plan{outer, extent, inner, start, step, length}, count, base, g), path) ==
944          Some{Pair.fst(F32, Array<F32> & Array<F32>, Slice.scattered((G, empty()),
945            Slice.target(Slice.Plan{outer, extent, inner, start, step, length}, Nat.add(Model.offset(depth, path), base))))} : Maybe<&2, F32>}> ->
946        {Model.observe(F32, depth, Slice.scatter_build(depth, splits, Slice.Plan{outer, extent, inner, start, step, length}, count, base, g), path) ==
947          Some{Pair.fst(F32, Array<F32> & Array<F32>, Slice.scattered((G, empty()),
948            Slice.target(Slice.Plan{outer, extent, inner, start, step, length}, Nat.add(Model.offset(depth, path), base))))} : Maybe<&2, F32>}}(
949        scatter_slot(depth, splits, Slice.Plan{outer, extent, inner, start, step, length}, count, base, dg, G, g, same, valid, path, inside)),
950      %Equal.sym(Maybe<&2, Nat>, Slice.target(Slice.Plan{outer, extent, inner, start, step, length}, Nat.add(Model.offset(depth, path), base)),
951        Some{Laws.position(length, inner, o, k, r)},
952        Equal.trans(Maybe<&2, Nat>, Slice.target(Slice.Plan{outer, extent, inner, start, step, length}, Nat.add(Model.offset(depth, path), base)),
953          Slice.target(Slice.Plan{outer, extent, inner, start, step, length}, Laws.position(extent, inner, o, Nat.add(start, Nat.mul(k, step)), r)),
954          Some{Laws.position(length, inner, o, k, r)},
955          Equal.cong(Nat, Maybe<&2, Nat>, x => Slice.target(Slice.Plan{outer, extent, inner, start, step, length}, x),
956            Nat.add(Model.offset(depth, path), base), Laws.position(extent, inner, o, Nat.add(start, Nat.mul(k, step)), r), at),
957          target_hit(outer, extent, inner, start, step, length, o, k, r, positive, k_below, i_below, r_below))) :
958        {Some{Pair.fst(F32, Array<F32> & Array<F32>, Slice.scattered((G, empty()), _))} ==
959          Model.observe(F32, dg, G, Refinement.route(dg, U32.from_nat(Laws.position(length, inner, o, k, r)))) : Maybe<&2, F32>}
960      {\{Parallel.Fact: q => q} :
961        Parallel.Fact<{Some{Pair.fst(F32, Array<F32> & Array<F32>, Slice.take(empty(), Array.get(F32, G, U32.from_nat(Laws.position(length, inner, o, k, r)))))} ==
962          Model.observe(F32, dg, G, Refinement.route(dg, U32.from_nat(Laws.position(length, inner, o, k, r)))) : Maybe<&2, F32>}> ->
963        {Some{Pair.fst(F32, Array<F32> & Array<F32>, Slice.take(empty(), Array.get(F32, G, U32.from_nat(Laws.position(length, inner, o, k, r)))))} ==
964          Model.observe(F32, dg, G, Refinement.route(dg, U32.from_nat(Laws.position(length, inner, o, k, r)))) : Maybe<&2, F32>}}(
965        read_word(dg, G, empty(), Laws.position(length, inner, o, k, r), valid)))}