src/tensor/batched_matmul_PROOF.bend

bendtorch-core@0.1.0.0 · 857 lines · 47 definitions

1import Base
2import bendtorch-core@0.1.0.0/src/shape/core.bend as Shape
3import bendtorch-core@0.1.0.0/src/storage/model.bend as Model
4import bendtorch-core@0.1.0.0/src/storage/refinement.bend as Refinement
5import bendtorch-core@0.1.0.0/src/tensor/core.bend as Core
6import bendtorch-core@0.1.0.0/src/tensor/lifecycle.bend as Lifecycle
7import bendtorch-core@0.1.0.0/src/tensor/matmul.bend as Matrix
8import bendtorch-core@0.1.0.0/src/tensor/batched_matmul.bend as Batched
9import bendtorch-core@0.1.0.0/src/tensor/batched_matmul_LAWS.bend as Laws
10import bendtorch-core@0.1.0.0/src/tensor/dot_fold.bend as DotFold
11import bendtorch-core@0.1.0.0/src/tensor/elementwise.bend as Elementwise
12import bendtorch-core@0.1.0.0/src/tensor/elementwise_parallel_proof.bend as Parallel
13import bendtorch-core@0.1.0.0/src/tensor/parallel.bend as Build
14import bendtorch-core@0.1.0.0/src/tensor/parallel_proofs.bend as Generic
15import bendtorch-core@0.1.0.0/src/tensor/matrix_index_proof.bend as MatrixIndex
16import bendtorch-core@0.1.0.0/src/shape/nat_proofs.bend as Nats
17import bendtorch-core@0.1.0.0/src/shape/index_proofs.bend as Index
18import bendtorch-core@0.1.0.0/src/storage/arith.bend as Arith
19
20# Which call answers a pair of operands.
21
22def Laws.batched_matmul_matrices(left, right, a, b, rows, inner, right_inner, columns, left_matrix, right_matrix):
23  %Equal.sym(List<&2, U32>, Shape.dimensions(left), [rows, inner], left_matrix) :
24    {Batched.routed(left, right, a, b, _, Shape.dimensions(right)) == Matrix.matmul(Core.Packed{left, a}, Core.Packed{right, b}) :
25      Result<&2, &1, Matrix.Error, Core.Dynamic & Core.Dynamic & Core.Dynamic>}
26  %Equal.sym(List<&2, U32>, Shape.dimensions(right), [right_inner, columns], right_matrix) :
27    {Batched.routed(left, right, a, b, [rows, inner], _) == Matrix.matmul(Core.Packed{left, a}, Core.Packed{right, b}) :
28      Result<&2, &1, Matrix.Error, Core.Dynamic & Core.Dynamic & Core.Dynamic>}
29  {==}
30
31def Laws.batched_matmul_left_scalar(left, right, a, b, scalar):
32  %Equal.sym(List<&2, U32>, Shape.dimensions(left), [], scalar) :
33    {Batched.routed(left, right, a, b, _, Shape.dimensions(right)) == Fail{Matrix.RankMismatch{}} :
34      Result<&2, &1, Matrix.Error, Core.Dynamic & Core.Dynamic & Core.Dynamic>}
35  {==}
36
37# A rank-0 right operand: whatever the left dimensions, the pair is not two
38# matrices, and the plan fails at one rank or the other.
39def right_scalar_routed(+left: Shape.Shape, +right: Shape.Shape, a: Core.Tensor<left>, b: Core.Tensor<right>, +dims: List<&2, U32>) ->
40  {Batched.routed(left, right, a, b, dims, []) == Fail{Matrix.RankMismatch{}} :
41    Result<&2, &1, Matrix.Error, Core.Dynamic & Core.Dynamic & Core.Dynamic>}:
42  match dims:
43    case Nil{}:
44      {==}
45    case +k <> Nil{}:
46      {==}
47    case +m <> +k <> Nil{}:
48      {==}
49    case +x <> +m <> +k <> Nil{}:
50      {==}
51    case +x <> +y <> +m <> +k <> Nil{}:
52      {==}
53    case +x <> +y <> +z <> +t <> +u <> rest:
54      {==}
55
56def Laws.batched_matmul_right_scalar(left, right, a, b, scalar):
57  %Equal.sym(List<&2, U32>, Shape.dimensions(right), [], scalar) :
58    {Batched.routed(left, right, a, b, Shape.dimensions(left), _) == Fail{Matrix.RankMismatch{}} :
59      Result<&2, &1, Matrix.Error, Core.Dynamic & Core.Dynamic & Core.Dynamic>}
60  right_scalar_routed(left, right, a, b, Shape.dimensions(left))
61
62def dot_compatible(+inner: U32, +right_inner: U32, mismatch: {U32.to_nat(inner) != U32.to_nat(right_inner) : Nat},
63  decision: Shape.BoolCase<Nat.is_eq(U32.to_nat(inner), U32.to_nat(right_inner))>) ->
64  {Batched.compatible(Batched.Operand{Batched.Batch{0n, 1, 1}, 1, inner, True{}}, Batched.Operand{Batched.Batch{0n, 1, 1}, 1, right_inner, True{}},
65    Nat.is_eq(U32.to_nat(inner), U32.to_nat(right_inner))) == Fail{Matrix.InnerDimensionMismatch{}} : Result<&2, &2, Matrix.Error, Batched.Plan>}:
66  match decision:
67    case Shape.BFalse{equal}:
68      %equal : {Batched.compatible(Batched.Operand{Batched.Batch{0n, 1, 1}, 1, inner, True{}}, Batched.Operand{Batched.Batch{0n, 1, 1}, 1, right_inner, True{}},
69        _) == Fail{Matrix.InnerDimensionMismatch{}} : Result<&2, &2, Matrix.Error, Batched.Plan>}
70      {==}
71    case Shape.BTrue{equal}:
72      Empty.absurd({Batched.compatible(Batched.Operand{Batched.Batch{0n, 1, 1}, 1, inner, True{}}, Batched.Operand{Batched.Batch{0n, 1, 1}, 1, right_inner, True{}},
73        Nat.is_eq(U32.to_nat(inner), U32.to_nat(right_inner))) == Fail{Matrix.InnerDimensionMismatch{}} : Result<&2, &2, Matrix.Error, Batched.Plan>},
74        mismatch(Lifecycle.natural_equal(U32.to_nat(inner), U32.to_nat(right_inner), equal)))
75
76def Laws.batched_matmul_dot_inner_error(left, right, a, b, inner, right_inner, left_vector, right_vector, mismatch):
77  %Equal.sym(List<&2, U32>, Shape.dimensions(left), [inner], left_vector) :
78    {Batched.routed(left, right, a, b, _, Shape.dimensions(right)) == Fail{Matrix.InnerDimensionMismatch{}} :
79      Result<&2, &1, Matrix.Error, Core.Dynamic & Core.Dynamic & Core.Dynamic>}
80  %Equal.sym(List<&2, U32>, Shape.dimensions(right), [right_inner], right_vector) :
81    {Batched.routed(left, right, a, b, [inner], _) == Fail{Matrix.InnerDimensionMismatch{}} :
82      Result<&2, &1, Matrix.Error, Core.Dynamic & Core.Dynamic & Core.Dynamic>}
83  %Equal.sym(Result<&2, &2, Matrix.Error, Batched.Plan>,
84    Batched.compatible(Batched.Operand{Batched.Batch{0n, 1, 1}, 1, inner, True{}}, Batched.Operand{Batched.Batch{0n, 1, 1}, 1, right_inner, True{}},
85      Nat.is_eq(U32.to_nat(inner), U32.to_nat(right_inner))), Fail{Matrix.InnerDimensionMismatch{}},
86    dot_compatible(inner, right_inner, mismatch, Shape.bool_case(Nat.is_eq(U32.to_nat(inner), U32.to_nat(right_inner))))) :
87    {Batched.planned(left, right, a, b, _) == Fail{Matrix.InnerDimensionMismatch{}} :
88      Result<&2, &1, Matrix.Error, Core.Dynamic & Core.Dynamic & Core.Dynamic>}
89  {==}
90
91# What every cell of the build holds.
92
93# A run of `Matrix.product_run` from offsets lp and rp: its total folds the
94# factor pairs it read, and both arrays come back unchanged.
95type StridedRun<-fuel: Nat, -dl: Nat, -dr: Nat, -A: Array<F32>, -B: Array<F32>, -ls: Nat, -rs: Nat, -lp: Nat, -rp: Nat,
96  -accumulator: F32, -result: F32 & (Array<F32> & Array<F32>)> is Data:
97  StridedRun{-values: List<&2, DotFold.Factors>, -logical: Laws.Strided(fuel, dl, dr, A, B, ls, rs, lp, rp, values),
98    -fold: {Pair.fst(F32, Array<F32> & Array<F32>, result) == DotFold.dot_fold(values, accumulator) : F32},
99    -left_kept: {Pair.fst(Array<F32>, Array<F32>, Pair.snd(F32, Array<F32> & Array<F32>, result)) == A : Array<F32>},
100    -right_kept: {Pair.snd(Array<F32>, Array<F32>, Pair.snd(F32, Array<F32> & Array<F32>, result)) == B : Array<F32>}}
101
102def run_extend(+rest: Nat, +dl: Nat, +dr: Nat, -A: Array<F32>, -B: Array<F32>, +ls: Nat, +rs: Nat, +lp: Nat, +rp: Nat,
103  -accumulator: F32, -x: F32, -y: F32,
104  -ex: {Model.observe(F32, dl, A, Refinement.route(dl, U32.from_nat(lp))) == Some{x} : Maybe<&2, F32>},
105  -ey: {Model.observe(F32, dr, B, Refinement.route(dr, U32.from_nat(rp))) == Some{y} : Maybe<&2, F32>},
106  -result: F32 & (Array<F32> & Array<F32>),
107  -run: Parallel.Fact<StridedRun<rest, dl, dr, A, B, ls, rs, Nat.add(ls, lp), Nat.add(rs, rp), F32.add(accumulator, F32.mul(x, y)), result>>) ->
108  Parallel.Fact<StridedRun<1n+rest, dl, dr, A, B, ls, rs, lp, rp, accumulator, result>>:
109  Parallel.Fact{{\{StridedRun: values => logical => fold => left_kept => right_kept =>
110    StridedRun{DotFold.Factors{x, y} <> values, Laws.StridedStep{ex, ey, logical}, fold, left_kept, right_kept}} :
111      StridedRun<rest, dl, dr, A, B, ls, rs, Nat.add(ls, lp), Nat.add(rs, rp), F32.add(accumulator, F32.mul(x, y)), result> ->
112      StridedRun<1n+rest, dl, dr, A, B, ls, rs, lp, rp, accumulator, result>}(
113    {\{Parallel.Fact: p => p} : Parallel.Fact<StridedRun<rest, dl, dr, A, B, ls, rs, Nat.add(ls, lp), Nat.add(rs, rp), F32.add(accumulator, F32.mul(x, y)), result>> ->
114      StridedRun<rest, dl, dr, A, B, ls, rs, Nat.add(ls, lp), Nat.add(rs, rp), F32.add(accumulator, F32.mul(x, y)), result>}(run))}
115
116# The reads at lp and rp hold the words of A and B, and keep both arrays.
117type Read<-dl: Nat, -dr: Nat, -A: Array<F32>, -B: Array<F32>, -lp: Nat, -rp: Nat, -pl: Array<F32> & F32, -pr: Array<F32> & F32> is Data:
118  Read{-ex: {Model.observe(F32, dl, A, Refinement.route(dl, U32.from_nat(lp))) == Some{Refinement.second(F32, pl)} : Maybe<&2, F32>},
119    -ey: {Model.observe(F32, dr, B, Refinement.route(dr, U32.from_nat(rp))) == Some{Refinement.second(F32, pr)} : Maybe<&2, F32>},
120    -kl: {Refinement.first(F32, pl) == A : Array<F32>}, -kr: {Refinement.first(F32, pr) == B : Array<F32>}}
121
122def reads(+dl: Nat, +dr: Nat, -A: Array<F32>, -B: Array<F32>, +lp: Nat, +rp: Nat,
123  -la: Array<F32>, -ra: Array<F32>, -el: {la == A : Array<F32>}, -er: {ra == B : Array<F32>},
124  -left_valid: Model.Balanced(F32, dl, A), -right_valid: Model.Balanced(F32, dr, B),
125  -pl: Array<F32> & F32, -exact_l: {pl == Array.get(F32, la, U32.from_nat(lp)) : Array<F32> & F32},
126  -pr: Array<F32> & F32, -exact_r: {pr == Array.get(F32, ra, U32.from_nat(rp)) : Array<F32> & F32}) ->
127  Parallel.Fact<Read<dl, dr, A, B, lp, rp, pl, pr>>:
128  Parallel.Fact{Read{
129    Equal.sym(Maybe<&2, F32>, Some{Refinement.second(F32, pl)}, Model.observe(F32, dl, A, Refinement.route(dl, U32.from_nat(lp))),
130      {\{Parallel.Fact: p => p} : Parallel.Fact<({Some{Refinement.second(F32, pl)} == Model.observe(F32, dl, A, Refinement.route(dl, U32.from_nat(lp))) : Maybe<&2, F32>})> ->
131        {Some{Refinement.second(F32, pl)} == Model.observe(F32, dl, A, Refinement.route(dl, U32.from_nat(lp))) : Maybe<&2, F32>}}(
132        Parallel.read_value(dl, la, A, U32.from_nat(lp), pl, exact_l, el, left_valid))),
133    Equal.sym(Maybe<&2, F32>, Some{Refinement.second(F32, pr)}, Model.observe(F32, dr, B, Refinement.route(dr, U32.from_nat(rp))),
134      {\{Parallel.Fact: p => p} : Parallel.Fact<({Some{Refinement.second(F32, pr)} == Model.observe(F32, dr, B, Refinement.route(dr, U32.from_nat(rp))) : Maybe<&2, F32>})> ->
135        {Some{Refinement.second(F32, pr)} == Model.observe(F32, dr, B, Refinement.route(dr, U32.from_nat(rp))) : Maybe<&2, F32>}}(
136        Parallel.read_value(dr, ra, B, U32.from_nat(rp), pr, exact_r, er, right_valid))),
137    {\{Parallel.Fact: p => p} : Parallel.Fact<({Refinement.first(F32, pl) == A : Array<F32>})> -> {Refinement.first(F32, pl) == A : Array<F32>}}(
138      Parallel.read_keeps(dl, la, A, U32.from_nat(lp), pl, exact_l, el, left_valid)),
139    {\{Parallel.Fact: p => p} : Parallel.Fact<({Refinement.first(F32, pr) == B : Array<F32>})> -> {Refinement.first(F32, pr) == B : Array<F32>}}(
140      Parallel.read_keeps(dr, ra, B, U32.from_nat(rp), pr, exact_r, er, right_valid))}}
141
142def run_last(+dl: Nat, +dr: Nat, -A: Array<F32>, -B: Array<F32>, +ls: Nat, +rs: Nat, +lp: Nat, +rp: Nat, -accumulator: F32,
143  pl: Array<F32> & F32, pr: Array<F32> & F32, -read: Read<dl, dr, A, B, lp, rp, pl, pr>) ->
144  Parallel.Fact<StridedRun<1n, dl, dr, A, B, ls, rs, lp, rp, accumulator, Matrix.product_run(0n, ls, rs, lp, rp, accumulator, pl, pr)>>:
145  (kl, x) = pl
146  (kr, y) = pr
147  -ex = {\{Read: ex => ey => left_kept => right_kept => ex} : Read<dl, dr, A, B, lp, rp, (kl, x), (kr, y)> ->
148    {Model.observe(F32, dl, A, Refinement.route(dl, U32.from_nat(lp))) == Some{x} : Maybe<&2, F32>}}(read)
149  -ey = {\{Read: ex => ey => left_kept => right_kept => ey} : Read<dl, dr, A, B, lp, rp, (kl, x), (kr, y)> ->
150    {Model.observe(F32, dr, B, Refinement.route(dr, U32.from_nat(rp))) == Some{y} : Maybe<&2, F32>}}(read)
151  -left_kept = {\{Read: ex => ey => left_kept => right_kept => left_kept} : Read<dl, dr, A, B, lp, rp, (kl, x), (kr, y)> ->
152    {kl == A : Array<F32>}}(read)
153  -right_kept = {\{Read: ex => ey => left_kept => right_kept => right_kept} : Read<dl, dr, A, B, lp, rp, (kl, x), (kr, y)> ->
154    {kr == B : Array<F32>}}(read)
155  run_extend(0n, dl, dr, A, B, ls, rs, lp, rp, accumulator, x, y, ex, ey, (F32.add(accumulator, F32.mul(x, y)), (kl, kr)),
156    Parallel.Fact{StridedRun{Nil{}, Unit{}, {==}, left_kept, right_kept}})
157
158def run_next(+rest: Nat, +dl: Nat, +dr: Nat, -A: Array<F32>, -B: Array<F32>, +ls: Nat, +rs: Nat, +lp: Nat, +rp: Nat, -accumulator: F32,
159  pl: Array<F32> & F32, pr: Array<F32> & F32, -read: Read<dl, dr, A, B, lp, rp, pl, pr>,
160  -next: @value: F32 -> @-l: Array<F32> -> @-r: Array<F32> -> @-el: {l == A : Array<F32>} -> @-er: {r == B : Array<F32>} ->
161    Parallel.Fact<StridedRun<1n+rest, dl, dr, A, B, ls, rs, Nat.add(ls, lp), Nat.add(rs, rp), value,
162      Matrix.product_run(rest, ls, rs, Nat.add(ls, lp), Nat.add(rs, rp), value,
163        Array.get(F32, l, U32.from_nat(Nat.add(ls, lp))), Array.get(F32, r, U32.from_nat(Nat.add(rs, rp))))>>) ->
164  Parallel.Fact<StridedRun<2n+rest, dl, dr, A, B, ls, rs, lp, rp, accumulator, Matrix.product_run(1n+rest, ls, rs, lp, rp, accumulator, pl, pr)>>:
165  (kl, x) = pl
166  (kr, y) = pr
167  -ex = {\{Read: ex => ey => left_kept => right_kept => ex} : Read<dl, dr, A, B, lp, rp, (kl, x), (kr, y)> ->
168    {Model.observe(F32, dl, A, Refinement.route(dl, U32.from_nat(lp))) == Some{x} : Maybe<&2, F32>}}(read)
169  -ey = {\{Read: ex => ey => left_kept => right_kept => ey} : Read<dl, dr, A, B, lp, rp, (kl, x), (kr, y)> ->
170    {Model.observe(F32, dr, B, Refinement.route(dr, U32.from_nat(rp))) == Some{y} : Maybe<&2, F32>}}(read)
171  -left_kept = {\{Read: ex => ey => left_kept => right_kept => left_kept} : Read<dl, dr, A, B, lp, rp, (kl, x), (kr, y)> ->
172    {kl == A : Array<F32>}}(read)
173  -right_kept = {\{Read: ex => ey => left_kept => right_kept => right_kept} : Read<dl, dr, A, B, lp, rp, (kl, x), (kr, y)> ->
174    {kr == B : Array<F32>}}(read)
175  run_extend(1n+rest, dl, dr, A, B, ls, rs, lp, rp, accumulator, x, y, ex, ey,
176    Matrix.product_run(rest, ls, rs, Nat.add(ls, lp), Nat.add(rs, rp), F32.add(accumulator, F32.mul(x, y)),
177      Array.get(F32, kl, U32.from_nat(Nat.add(ls, lp))), Array.get(F32, kr, U32.from_nat(Nat.add(rs, rp)))),
178    next(F32.add(accumulator, F32.mul(x, y)))(kl)(kr)(left_kept)(right_kept))
179
180# A run of fuel + 1 terms from lp and rp folds the factor pairs it read and
181# hands both arrays back.
182def run_eq(+fuel: Nat, +dl: Nat, +dr: Nat, -A: Array<F32>, -B: Array<F32>, +ls: Nat, +rs: Nat, +lp: Nat, +rp: Nat, -accumulator: F32,
183  -la: Array<F32>, -ra: Array<F32>, -el: {la == A : Array<F32>}, -er: {ra == B : Array<F32>},
184  -left_valid: Model.Balanced(F32, dl, A), -right_valid: Model.Balanced(F32, dr, B),
185  pl: Array<F32> & F32, -exact_l: {pl == Array.get(F32, la, U32.from_nat(lp)) : Array<F32> & F32},
186  pr: Array<F32> & F32, -exact_r: {pr == Array.get(F32, ra, U32.from_nat(rp)) : Array<F32> & F32}) ->
187  Parallel.Fact<StridedRun<1n+fuel, dl, dr, A, B, ls, rs, lp, rp, accumulator, Matrix.product_run(fuel, ls, rs, lp, rp, accumulator, pl, pr)>>:
188  match fuel:
189    case 0n:
190      run_last(dl, dr, A, B, ls, rs, lp, rp, accumulator, pl, pr,
191        {\{Parallel.Fact: p => p} : Parallel.Fact<Read<dl, dr, A, B, lp, rp, pl, pr>> -> Read<dl, dr, A, B, lp, rp, pl, pr>}(
192          reads(dl, dr, A, B, lp, rp, la, ra, el, er, left_valid, right_valid, pl, exact_l, pr, exact_r)))
193    case 1n+rest:
194      run_next(rest, dl, dr, A, B, ls, rs, lp, rp, accumulator, pl, pr,
195        {\{Parallel.Fact: p => p} : Parallel.Fact<Read<dl, dr, A, B, lp, rp, pl, pr>> -> Read<dl, dr, A, B, lp, rp, pl, pr>}(
196          reads(dl, dr, A, B, lp, rp, la, ra, el, er, left_valid, right_valid, pl, exact_l, pr, exact_r)),
197        value => l => r => el => er => run_eq(rest, dl, dr, A, B, ls, rs, Nat.add(ls, lp), Nat.add(rs, rp), value, l, r, el, er,
198          left_valid, right_valid, Array.get(F32, l, U32.from_nat(Nat.add(ls, lp))), {==}, Array.get(F32, r, U32.from_nat(Nat.add(rs, rp))), {==}))
199
200# `Batched.run_from` of `fuel` terms from +0.0.
201def from_run(fuel: Nat, +dl: Nat, +dr: Nat, -A: Array<F32>, -B: Array<F32>, +ls: Nat, +rs: Nat, +lp: Nat, +rp: Nat,
202  la: Array<F32>, ra: Array<F32>, -el: {la == A : Array<F32>}, -er: {ra == B : Array<F32>},
203  -left_valid: Model.Balanced(F32, dl, A), -right_valid: Model.Balanced(F32, dr, B)) ->
204  Parallel.Fact<StridedRun<fuel, dl, dr, A, B, ls, rs, lp, rp, 0.0, Batched.run_from(fuel, ls, rs, lp, rp, la, ra)>>:
205  match fuel:
206    case 0n:
207      Parallel.Fact{StridedRun{Nil{}, Unit{}, {==}, el, er}}
208    case 1n+rest:
209      run_eq(rest, dl, dr, A, B, ls, rs, lp, rp, 0.0, la, ra, el, er, left_valid, right_valid,
210        Array.get(F32, la, U32.from_nat(lp)), {==}, Array.get(F32, ra, U32.from_nat(rp)), {==})
211
212# The batched cell function of `Parallel.build`, from the input arrays.
213def cell_of(+lo: Matrix.Order, +ro: Matrix.Order, +inner: Nat, +rows: Nat, +columns: Nat, +l1: Nat, +l2: Nat, +r1: Nat, +r2: Nat, +o2: Nat,
214  +position: Nat, A: Array<F32>, B: Array<F32>, inside: Bool) -> F32 & (Array<F32> & Array<F32>):
215  Batched.cell_at(Batched.layout(Laws.layout_of(lo, ro, inner, rows, columns, l1, l2, r1, r2, o2)), position, (A, B), inside)
216
217def dot_run(+lo: Matrix.Order, +ro: Matrix.Order, +inner: Nat, +rows: Nat, +columns: Nat, +l1: Nat, +l2: Nat, +r1: Nat, +r2: Nat, +o2: Nat,
218  +position: Nat, +dl: Nat, +dr: Nat, A: Array<F32>, B: Array<F32>,
219  -left_valid: Model.Balanced(F32, dl, A), -right_valid: Model.Balanced(F32, dr, B)) ->
220  Parallel.Fact<StridedRun<inner, dl, dr, A, B, Matrix.stride(lo, rows), Matrix.stride(ro, columns),
221    Batched.left_start(Laws.layout_of(lo, ro, inner, rows, columns, l1, l2, r1, r2, o2), position),
222    Batched.right_start(Laws.layout_of(lo, ro, inner, rows, columns, l1, l2, r1, r2, o2), position), 0.0,
223    cell_of(lo, ro, inner, rows, columns, l1, l2, r1, r2, o2, position, A, B, True{})>>:
224  from_run(inner, dl, dr, A, B, Matrix.stride(lo, rows), Matrix.stride(ro, columns),
225    Batched.left_start(Laws.layout_of(lo, ro, inner, rows, columns, l1, l2, r1, r2, o2), position),
226    Batched.right_start(Laws.layout_of(lo, ro, inner, rows, columns, l1, l2, r1, r2, o2), position), A, B, {==}, {==}, left_valid, right_valid)
227
228def eta(-X: Type, -Y: Type, pair: X & Y) -> {pair == (Pair.fst(X, Y, pair), Pair.snd(X, Y, pair)) : X & Y}:
229  (x, y) = pair
230  {==}
231
232# Every cell hands both input arrays back.
233def cell_keeps(+lo: Matrix.Order, +ro: Matrix.Order, +inner: Nat, +rows: Nat, +columns: Nat, +l1: Nat, +l2: Nat, +r1: Nat, +r2: Nat, +o2: Nat,
234  +position: Nat, +dl: Nat, +dr: Nat, A: Array<F32>, B: Array<F32>, inside: Bool,
235  -left_valid: Model.Balanced(F32, dl, A), -right_valid: Model.Balanced(F32, dr, B)) ->
236  Parallel.Fact<{Pair.snd(F32, Array<F32> & Array<F32>, cell_of(lo, ro, inner, rows, columns, l1, l2, r1, r2, o2, position, A, B, inside)) == (A, B) : Array<F32> & Array<F32>}>:
237  match inside:
238    case False{}:
239      Parallel.Fact{{==}}
240    case True{}:
241      Parallel.Fact{{\{StridedRun: values => logical => fold => left_kept => right_kept =>
242        Equal.trans(Array<F32> & Array<F32>, Pair.snd(F32, Array<F32> & Array<F32>, cell_of(lo, ro, inner, rows, columns, l1, l2, r1, r2, o2, position, A, B, True{})),
243          (Pair.fst(Array<F32>, Array<F32>, Pair.snd(F32, Array<F32> & Array<F32>, cell_of(lo, ro, inner, rows, columns, l1, l2, r1, r2, o2, position, A, B, True{}))),
244            Pair.snd(Array<F32>, Array<F32>, Pair.snd(F32, Array<F32> & Array<F32>, cell_of(lo, ro, inner, rows, columns, l1, l2, r1, r2, o2, position, A, B, True{})))),
245          (A, B),
246          eta(Array<F32>, Array<F32>, Pair.snd(F32, Array<F32> & Array<F32>, cell_of(lo, ro, inner, rows, columns, l1, l2, r1, r2, o2, position, A, B, True{}))),
247          Parallel.pair_equal(Pair.fst(Array<F32>, Array<F32>, Pair.snd(F32, Array<F32> & Array<F32>, cell_of(lo, ro, inner, rows, columns, l1, l2, r1, r2, o2, position, A, B, True{}))),
248            Pair.snd(Array<F32>, Array<F32>, Pair.snd(F32, Array<F32> & Array<F32>, cell_of(lo, ro, inner, rows, columns, l1, l2, r1, r2, o2, position, A, B, True{}))),
249            A, B, left_kept, right_kept))} :
250        StridedRun<inner, dl, dr, A, B, Matrix.stride(lo, rows), Matrix.stride(ro, columns),
251          Batched.left_start(Laws.layout_of(lo, ro, inner, rows, columns, l1, l2, r1, r2, o2), position),
252          Batched.right_start(Laws.layout_of(lo, ro, inner, rows, columns, l1, l2, r1, r2, o2), position), 0.0,
253          cell_of(lo, ro, inner, rows, columns, l1, l2, r1, r2, o2, position, A, B, True{})> ->
254        {Pair.snd(F32, Array<F32> & Array<F32>, cell_of(lo, ro, inner, rows, columns, l1, l2, r1, r2, o2, position, A, B, True{})) == (A, B) : Array<F32> & Array<F32>}}(
255        {\{Parallel.Fact: p => p} :
256          Parallel.Fact<StridedRun<inner, dl, dr, A, B, Matrix.stride(lo, rows), Matrix.stride(ro, columns),
257            Batched.left_start(Laws.layout_of(lo, ro, inner, rows, columns, l1, l2, r1, r2, o2), position),
258            Batched.right_start(Laws.layout_of(lo, ro, inner, rows, columns, l1, l2, r1, r2, o2), position), 0.0,
259            cell_of(lo, ro, inner, rows, columns, l1, l2, r1, r2, o2, position, A, B, True{})>> ->
260          StridedRun<inner, dl, dr, A, B, Matrix.stride(lo, rows), Matrix.stride(ro, columns),
261            Batched.left_start(Laws.layout_of(lo, ro, inner, rows, columns, l1, l2, r1, r2, o2), position),
262            Batched.right_start(Laws.layout_of(lo, ro, inner, rows, columns, l1, l2, r1, r2, o2), position), 0.0,
263            cell_of(lo, ro, inner, rows, columns, l1, l2, r1, r2, o2, position, A, B, True{})>}(
264          dot_run(lo, ro, inner, rows, columns, l1, l2, r1, r2, o2, position, dl, dr, A, B, left_valid, right_valid)))}
265
266# What the cell function computes is what CellSpec names.
267def cell_spec(+lo: Matrix.Order, +ro: Matrix.Order, +inner: Nat, +rows: Nat, +columns: Nat, +l1: Nat, +l2: Nat, +r1: Nat, +r2: Nat, +o2: Nat,
268  +position: Nat, +dl: Nat, +dr: Nat, A: Array<F32>, B: Array<F32>, inside: Bool, -observed: Maybe<&2, F32>,
269  -computed: {observed == Some{Pair.fst(F32, Array<F32> & Array<F32>, cell_of(lo, ro, inner, rows, columns, l1, l2, r1, r2, o2, position, A, B, inside))} : Maybe<&2, F32>},
270  -left_valid: Model.Balanced(F32, dl, A), -right_valid: Model.Balanced(F32, dr, B)) ->
271  Parallel.Fact<Laws.CellSpec(dl, dr, A, B, lo, ro, inner, rows, columns, l1, l2, r1, r2, o2, position, inside, observed)>:
272  match inside:
273    case False{}:
274      Parallel.Fact{computed}
275    case True{}:
276      Parallel.Fact{{\{StridedRun: values => logical => fold => left_kept => right_kept =>
277        Laws.CellAt{Pair.fst(F32, Array<F32> & Array<F32>, cell_of(lo, ro, inner, rows, columns, l1, l2, r1, r2, o2, position, A, B, True{})),
278          computed, values, logical, fold}} :
279        StridedRun<inner, dl, dr, A, B, Matrix.stride(lo, rows), Matrix.stride(ro, columns),
280          Batched.left_start(Laws.layout_of(lo, ro, inner, rows, columns, l1, l2, r1, r2, o2), position),
281          Batched.right_start(Laws.layout_of(lo, ro, inner, rows, columns, l1, l2, r1, r2, o2), position), 0.0,
282          cell_of(lo, ro, inner, rows, columns, l1, l2, r1, r2, o2, position, A, B, True{})> ->
283        Laws.CellAt<dl, dr, A, B, lo, ro, inner, rows, columns, l1, l2, r1, r2, o2, position, observed>}(
284        {\{Parallel.Fact: p => p} :
285          Parallel.Fact<StridedRun<inner, dl, dr, A, B, Matrix.stride(lo, rows), Matrix.stride(ro, columns),
286            Batched.left_start(Laws.layout_of(lo, ro, inner, rows, columns, l1, l2, r1, r2, o2), position),
287            Batched.right_start(Laws.layout_of(lo, ro, inner, rows, columns, l1, l2, r1, r2, o2), position), 0.0,
288            cell_of(lo, ro, inner, rows, columns, l1, l2, r1, r2, o2, position, A, B, True{})>> ->
289          StridedRun<inner, dl, dr, A, B, Matrix.stride(lo, rows), Matrix.stride(ro, columns),
290            Batched.left_start(Laws.layout_of(lo, ro, inner, rows, columns, l1, l2, r1, r2, o2), position),
291            Batched.right_start(Laws.layout_of(lo, ro, inner, rows, columns, l1, l2, r1, r2, o2), position), 0.0,
292            cell_of(lo, ro, inner, rows, columns, l1, l2, r1, r2, o2, position, A, B, True{})>}(
293          dot_run(lo, ro, inner, rows, columns, l1, l2, r1, r2, o2, position, dl, dr, A, B, left_valid, right_valid)))}
294
295def Laws.batched_build_cells(depth, splits, dl, dr, A, B, lo, ro, inner, rows, columns, l1, l2, r1, r2, o2, count, base, la, ra, el, er,
296  left_valid, right_valid):
297  -split = {\{Parallel.Fact: q => q} : Parallel.Fact<({Pair.fst(Array<F32> & Array<F32>, Array<F32> & Array<F32>, Elementwise.split((A, B))) == (A, B) : Array<F32> & Array<F32>} &
298    {Pair.snd(Array<F32> & Array<F32>, Array<F32> & Array<F32>, Elementwise.split((A, B))) == (A, B) : Array<F32> & Array<F32>})> ->
299    {Pair.fst(Array<F32> & Array<F32>, Array<F32> & Array<F32>, Elementwise.split((A, B))) == (A, B) : Array<F32> & Array<F32>} &
300    {Pair.snd(Array<F32> & Array<F32>, Array<F32> & Array<F32>, Elementwise.split((A, B))) == (A, B) : Array<F32> & Array<F32>}}(
301    Parallel.split_equal(dl, dr, A, B, left_valid, right_valid))
302  -cells = {\{Generic.Fact: q => q} :
303    Generic.Fact<Generic.Cells(~Batched.Layout, ~Array<F32>, ~Array<F32>, ~Batched.cell_at, Batched.layout(Laws.layout_of(lo, ro, inner, rows, columns, l1, l2, r1, r2, o2)),
304      (A, B), depth, count, base, Batched.build(depth, splits, Laws.layout_of(lo, ro, inner, rows, columns, l1, l2, r1, r2, o2), count, base, la, ra))> ->
305    Generic.Cells(~Batched.Layout, ~Array<F32>, ~Array<F32>, ~Batched.cell_at, Batched.layout(Laws.layout_of(lo, ro, inner, rows, columns, l1, l2, r1, r2, o2)),
306      (A, B), depth, count, base, Batched.build(depth, splits, Laws.layout_of(lo, ro, inner, rows, columns, l1, l2, r1, r2, o2), count, base, la, ra))}(
307    Generic.build_cells(~Batched.Layout, ~Batched.Layout, ~Array<F32>, ~Array<F32>, ~Batched.layout, ~Elementwise.split, ~Batched.cell_at,
308      Laws.layout_of(lo, ro, inner, rows, columns, l1, l2, r1, r2, o2), (A, B), depth, splits, count, base, (la, ra), Parallel.pair_equal(la, ra, A, B, el, er),
309      at => inside => {\{Parallel.Fact: q => q} :
310        Parallel.Fact<({Pair.snd(F32, Array<F32> & Array<F32>, cell_of(lo, ro, inner, rows, columns, l1, l2, r1, r2, o2, at, A, B, inside)) == (A, B) : Array<F32> & Array<F32>})> ->
311        {Pair.snd(F32, Array<F32> & Array<F32>, cell_of(lo, ro, inner, rows, columns, l1, l2, r1, r2, o2, at, A, B, inside)) == (A, B) : Array<F32> & Array<F32>}}(
312        cell_keeps(lo, ro, inner, rows, columns, l1, l2, r1, r2, o2, at, dl, dr, A, B, inside, left_valid, right_valid)),
313      Pair.fst({Pair.fst(Array<F32> & Array<F32>, Array<F32> & Array<F32>, Elementwise.split((A, B))) == (A, B) : Array<F32> & Array<F32>},
314        {Pair.snd(Array<F32> & Array<F32>, Array<F32> & Array<F32>, Elementwise.split((A, B))) == (A, B) : Array<F32> & Array<F32>}, split),
315      Pair.snd({Pair.fst(Array<F32> & Array<F32>, Array<F32> & Array<F32>, Elementwise.split((A, B))) == (A, B) : Array<F32> & Array<F32>},
316        {Pair.snd(Array<F32> & Array<F32>, Array<F32> & Array<F32>, Elementwise.split((A, B))) == (A, B) : Array<F32> & Array<F32>}, split)))
317  Laws.Holds{path => {\{Parallel.Fact: q => q} :
318    Parallel.Fact<Laws.CellSpec(dl, dr, A, B, lo, ro, inner, rows, columns, l1, l2, r1, r2, o2, Nat.add(Model.offset(depth, path), base),
319      Nat.is_lt(Nat.add(Model.offset(depth, path), base), count),
320      Model.observe(F32, depth, Batched.build(depth, splits, Laws.layout_of(lo, ro, inner, rows, columns, l1, l2, r1, r2, o2), count, base, la, ra), path))> ->
321    Laws.CellSpec(dl, dr, A, B, lo, ro, inner, rows, columns, l1, l2, r1, r2, o2, Nat.add(Model.offset(depth, path), base),
322      Nat.is_lt(Nat.add(Model.offset(depth, path), base), count),
323      Model.observe(F32, depth, Batched.build(depth, splits, Laws.layout_of(lo, ro, inner, rows, columns, l1, l2, r1, r2, o2), count, base, la, ra), path))}(
324    cell_spec(lo, ro, inner, rows, columns, l1, l2, r1, r2, o2, Nat.add(Model.offset(depth, path), base), dl, dr, A, B,
325      Nat.is_lt(Nat.add(Model.offset(depth, path), base), count),
326      Model.observe(F32, depth, Batched.build(depth, splits, Laws.layout_of(lo, ro, inner, rows, columns, l1, l2, r1, r2, o2), count, base, la, ra), path),
327      cells(path), left_valid, right_valid))}
328
329def Laws.batched_build_balanced(depth, splits, given, count, base, la, ra):
330  Batched.build_balanced(depth, splits, given, count, base, la, ra)
331
332# The cells of the build in the output's coordinates.
333
334# Division and remainder of q * b + r by b, for r below b.
335def div_exact(+q: Nat, +b: Nat, +r: Nat, below: {Nat.is_lt(r, b) == True{} : Bool}) -> {Nat.div(Nat.add(Nat.mul(q, b), r), b) == q : Nat}:
336  Equal.cong(Nat & Nat, Nat, pair => Pair.fst(Nat, Nat, pair), Nat.divmod(Nat.add(Nat.mul(q, b), r), b), (q, r),
337    Equal.sym(Nat & Nat, (q, r), Nat.divmod(Nat.add(Nat.mul(q, b), r), b), Nats.divmod_exact(q, b, r, MatrixIndex.lt_from(r, b, below))))
338
339def mod_exact(+q: Nat, +b: Nat, +r: Nat, below: {Nat.is_lt(r, b) == True{} : Bool}) -> {Nat.mod(Nat.add(Nat.mul(q, b), r), b) == r : Nat}:
340  Equal.cong(Nat & Nat, Nat, pair => Pair.snd(Nat, Nat, pair), Nat.divmod(Nat.add(Nat.mul(q, b), r), b), (q, r),
341    Equal.sym(Nat & Nat, (q, r), Nat.divmod(Nat.add(Nat.mul(q, b), r), b), Nats.divmod_exact(q, b, r, MatrixIndex.lt_from(r, b, below))))
342
343# The batch number of the build reads each operand's batch as the shape
344# module broadcasts it.
345# An index of a batch axis: the stretch of the build, the shape module's
346# broadcast index.
347def stretch_spec(+extent: Nat, +index: Nat) -> {Batched.stretch(index, Nat.is_eq(extent, 1n)) == Shape.broadcast_index(extent, index) : Nat}:
348  match extent:
349    case 0n:
350      {==}
351    case 1n+e:
352      match e:
353        case 0n:
354          {==}
355        case 1n+f:
356          {==}
357
358# The batch number of the build reads each operand's batch as the shape
359# module broadcasts it.
360def batch_number_spec(+first: Nat, +second: Nat, +b1: Nat, +b2: Nat) ->
361  {Batched.batch_number(first, second, b1, b2) == Nat.add(Nat.mul(Shape.broadcast_index(first, b1), second), Shape.broadcast_index(second, b2)) : Nat}:
362  %Equal.sym(Nat, Batched.stretch(b1, Nat.is_eq(first, 1n)), Shape.broadcast_index(first, b1), stretch_spec(first, b1)) :
363    {Nat.add(Nat.mul(_, second), Batched.stretch(b2, Nat.is_eq(second, 1n))) ==
364      Nat.add(Nat.mul(Shape.broadcast_index(first, b1), second), Shape.broadcast_index(second, b2)) : Nat}
365  %Equal.sym(Nat, Batched.stretch(b2, Nat.is_eq(second, 1n)), Shape.broadcast_index(second, b2), stretch_spec(second, b2)) :
366    {Nat.add(Nat.mul(Shape.broadcast_index(first, b1), second), _) ==
367      Nat.add(Nat.mul(Shape.broadcast_index(first, b1), second), Shape.broadcast_index(second, b2)) : Nat}
368  {==}
369
370def read_offset_spec(+order: Matrix.Order, +inner: Nat, +outers: Nat, +outer: Nat, +k: Nat) ->
371  {Matrix.read_offset(order, inner, outers, outer, k) == Laws.matrix_offset(order, inner, outers, outer, k) : Nat}:
372  match order:
373    case Matrix.Kept{}:
374      {==}
375    case Matrix.Swapped{}:
376      {==}
377
378# The first left read of a cell, from its batch number `batch` and its
379# offset `within` inside its matrix.
380def left_from(+lo: Matrix.Order, +inner: Nat, +rows: Nat, +columns: Nat, +l1: Nat, +l2: Nat, +o2: Nat, +batch: Nat, +within: Nat) -> Nat:
381  Nat.add(Nat.mul(Batched.batch_number(l1, l2, Nat.div(batch, o2), Nat.mod(batch, o2)), Nat.mul(rows, inner)),
382    Matrix.read_offset(lo, inner, rows, Nat.div(within, columns), 0n))
383
384def right_from(+ro: Matrix.Order, +inner: Nat, +rows: Nat, +columns: Nat, +r1: Nat, +r2: Nat, +o2: Nat, +batch: Nat, +within: Nat) -> Nat:
385  Nat.add(Nat.mul(Batched.batch_number(r1, r2, Nat.div(batch, o2), Nat.mod(batch, o2)), Nat.mul(inner, columns)),
386    Matrix.read_offset(ro, inner, columns, Nat.mod(within, columns), 0n))
387
388def left_start_spec(+lo: Matrix.Order, +ro: Matrix.Order, +inner: Nat, +rows: Nat, +columns: Nat, +l1: Nat, +l2: Nat, +r1: Nat, +r2: Nat, +o2: Nat,
389  +b1: Nat, +b2: Nat, +row: Nat, +column: Nat,
390  +row_below: {Nat.is_lt(row, rows) == True{} : Bool}, +column_below: {Nat.is_lt(column, columns) == True{} : Bool},
391  +batch_below: {Nat.is_lt(b2, o2) == True{} : Bool}) ->
392  {Batched.left_start(Laws.layout_of(lo, ro, inner, rows, columns, l1, l2, r1, r2, o2), Laws.cell_position(o2, rows, columns, b1, b2, row, column)) ==
393    Laws.operand_offset(lo, inner, rows, l1, l2, b1, b2, row, 0n) : Nat}:
394  +q = Nat.add(Nat.mul(b1, o2), b2)
395  +r = Nat.add(Nat.mul(row, columns), column)
396  +rc = Nat.mul(rows, columns)
397  +within = MatrixIndex.lt_mul_add(row, rows, column, columns, row_below, column_below)
398  %Equal.sym(Nat, Nat.div(Nat.add(Nat.mul(q, rc), r), rc), q, div_exact(q, rc, r, within)) :
399    {left_from(lo, inner, rows, columns, l1, l2, o2, _, Nat.mod(Nat.add(Nat.mul(q, rc), r), rc)) == Laws.operand_offset(lo, inner, rows, l1, l2, b1, b2, row, 0n) : Nat}
400  %Equal.sym(Nat, Nat.mod(Nat.add(Nat.mul(q, rc), r), rc), r, mod_exact(q, rc, r, within)) :
401    {left_from(lo, inner, rows, columns, l1, l2, o2, q, _) == Laws.operand_offset(lo, inner, rows, l1, l2, b1, b2, row, 0n) : Nat}
402  %Equal.sym(Nat, Nat.div(q, o2), b1, div_exact(b1, o2, b2, batch_below)) :
403    {Nat.add(Nat.mul(Batched.batch_number(l1, l2, _, Nat.mod(q, o2)), Nat.mul(rows, inner)), Matrix.read_offset(lo, inner, rows, Nat.div(r, columns), 0n)) ==
404      Laws.operand_offset(lo, inner, rows, l1, l2, b1, b2, row, 0n) : Nat}
405  %Equal.sym(Nat, Nat.mod(q, o2), b2, mod_exact(b1, o2, b2, batch_below)) :
406    {Nat.add(Nat.mul(Batched.batch_number(l1, l2, b1, _), Nat.mul(rows, inner)), Matrix.read_offset(lo, inner, rows, Nat.div(r, columns), 0n)) ==
407      Laws.operand_offset(lo, inner, rows, l1, l2, b1, b2, row, 0n) : Nat}
408  %Equal.sym(Nat, Nat.div(r, columns), row, div_exact(row, columns, column, column_below)) :
409    {Nat.add(Nat.mul(Batched.batch_number(l1, l2, b1, b2), Nat.mul(rows, inner)), Matrix.read_offset(lo, inner, rows, _, 0n)) ==
410      Laws.operand_offset(lo, inner, rows, l1, l2, b1, b2, row, 0n) : Nat}
411  %Equal.sym(Nat, Batched.batch_number(l1, l2, b1, b2), Nat.add(Nat.mul(Shape.broadcast_index(l1, b1), l2), Shape.broadcast_index(l2, b2)),
412    batch_number_spec(l1, l2, b1, b2)) :
413    {Nat.add(Nat.mul(_, Nat.mul(rows, inner)), Matrix.read_offset(lo, inner, rows, row, 0n)) ==
414      Laws.operand_offset(lo, inner, rows, l1, l2, b1, b2, row, 0n) : Nat}
415  %Equal.sym(Nat, Matrix.read_offset(lo, inner, rows, row, 0n), Laws.matrix_offset(lo, inner, rows, row, 0n), read_offset_spec(lo, inner, rows, row, 0n)) :
416    {Nat.add(Nat.mul(Nat.add(Nat.mul(Shape.broadcast_index(l1, b1), l2), Shape.broadcast_index(l2, b2)), Nat.mul(rows, inner)), _) ==
417      Laws.operand_offset(lo, inner, rows, l1, l2, b1, b2, row, 0n) : Nat}
418  {==}
419
420def right_start_spec(+lo: Matrix.Order, +ro: Matrix.Order, +inner: Nat, +rows: Nat, +columns: Nat, +l1: Nat, +l2: Nat, +r1: Nat, +r2: Nat, +o2: Nat,
421  +b1: Nat, +b2: Nat, +row: Nat, +column: Nat,
422  +row_below: {Nat.is_lt(row, rows) == True{} : Bool}, +column_below: {Nat.is_lt(column, columns) == True{} : Bool},
423  +batch_below: {Nat.is_lt(b2, o2) == True{} : Bool}) ->
424  {Batched.right_start(Laws.layout_of(lo, ro, inner, rows, columns, l1, l2, r1, r2, o2), Laws.cell_position(o2, rows, columns, b1, b2, row, column)) ==
425    Laws.operand_offset(ro, inner, columns, r1, r2, b1, b2, column, 0n) : Nat}:
426  +q = Nat.add(Nat.mul(b1, o2), b2)
427  +r = Nat.add(Nat.mul(row, columns), column)
428  +rc = Nat.mul(rows, columns)
429  +within = MatrixIndex.lt_mul_add(row, rows, column, columns, row_below, column_below)
430  %Equal.sym(Nat, Nat.div(Nat.add(Nat.mul(q, rc), r), rc), q, div_exact(q, rc, r, within)) :
431    {right_from(ro, inner, rows, columns, r1, r2, o2, _, Nat.mod(Nat.add(Nat.mul(q, rc), r), rc)) == Laws.operand_offset(ro, inner, columns, r1, r2, b1, b2, column, 0n) : Nat}
432  %Equal.sym(Nat, Nat.mod(Nat.add(Nat.mul(q, rc), r), rc), r, mod_exact(q, rc, r, within)) :
433    {right_from(ro, inner, rows, columns, r1, r2, o2, q, _) == Laws.operand_offset(ro, inner, columns, r1, r2, b1, b2, column, 0n) : Nat}
434  %Equal.sym(Nat, Nat.div(q, o2), b1, div_exact(b1, o2, b2, batch_below)) :
435    {Nat.add(Nat.mul(Batched.batch_number(r1, r2, _, Nat.mod(q, o2)), Nat.mul(inner, columns)), Matrix.read_offset(ro, inner, columns, Nat.mod(r, columns), 0n)) ==
436      Laws.operand_offset(ro, inner, columns, r1, r2, b1, b2, column, 0n) : Nat}
437  %Equal.sym(Nat, Nat.mod(q, o2), b2, mod_exact(b1, o2, b2, batch_below)) :
438    {Nat.add(Nat.mul(Batched.batch_number(r1, r2, b1, _), Nat.mul(inner, columns)), Matrix.read_offset(ro, inner, columns, Nat.mod(r, columns), 0n)) ==
439      Laws.operand_offset(ro, inner, columns, r1, r2, b1, b2, column, 0n) : Nat}
440  %Equal.sym(Nat, Nat.mod(r, columns), column, mod_exact(row, columns, column, column_below)) :
441    {Nat.add(Nat.mul(Batched.batch_number(r1, r2, b1, b2), Nat.mul(inner, columns)), Matrix.read_offset(ro, inner, columns, _, 0n)) ==
442      Laws.operand_offset(ro, inner, columns, r1, r2, b1, b2, column, 0n) : Nat}
443  %Equal.sym(Nat, Batched.batch_number(r1, r2, b1, b2), Nat.add(Nat.mul(Shape.broadcast_index(r1, b1), r2), Shape.broadcast_index(r2, b2)),
444    batch_number_spec(r1, r2, b1, b2)) :
445    {Nat.add(Nat.mul(_, Nat.mul(inner, columns)), Matrix.read_offset(ro, inner, columns, column, 0n)) ==
446      Laws.operand_offset(ro, inner, columns, r1, r2, b1, b2, column, 0n) : Nat}
447  %Equal.sym(Nat, Matrix.read_offset(ro, inner, columns, column, 0n), Laws.matrix_offset(ro, inner, columns, column, 0n), read_offset_spec(ro, inner, columns, column, 0n)) :
448    {Nat.add(Nat.mul(Nat.add(Nat.mul(Shape.broadcast_index(r1, b1), r2), Shape.broadcast_index(r2, b2)), Nat.mul(inner, columns)), _) ==
449      Laws.operand_offset(ro, inner, columns, r1, r2, b1, b2, column, 0n) : Nat}
450  %Equal.sym(Nat, Nat.mul(inner, columns), Nat.mul(columns, inner), Nats.mul_comm(inner, columns)) :
451    {Nat.add(Nat.mul(Nat.add(Nat.mul(Shape.broadcast_index(r1, b1), r2), Shape.broadcast_index(r2, b2)), _), Laws.matrix_offset(ro, inner, columns, column, 0n)) ==
452      Laws.operand_offset(ro, inner, columns, r1, r2, b1, b2, column, 0n) : Nat}
453  {==}
454
455def batch_base(+inner: Nat, +outers: Nat, +e1: Nat, +e2: Nat, +b1: Nat, +b2: Nat) -> Nat:
456  Nat.mul(Nat.add(Nat.mul(Shape.broadcast_index(e1, b1), e2), Shape.broadcast_index(e2, b2)), Nat.mul(outers, inner))
457
458# Term k + 1 of an operand sits one stride after term k.
459def operand_step(+order: Matrix.Order, +inner: Nat, +outers: Nat, +e1: Nat, +e2: Nat, +b1: Nat, +b2: Nat, +outer: Nat, +k: Nat) ->
460  {Nat.add(Matrix.stride(order, outers), Laws.operand_offset(order, inner, outers, e1, e2, b1, b2, outer, k)) ==
461    Laws.operand_offset(order, inner, outers, e1, e2, b1, b2, outer, 1n+k) : Nat}:
462  match order:
463    case Matrix.Kept{}:
464      %Equal.sym(Nat, Nat.add(Nat.mul(outer, inner), 1n+k), 1n+Nat.add(Nat.mul(outer, inner), k), Arith.add_succ(Nat.mul(outer, inner), k)) :
465        {1n+Nat.add(batch_base(inner, outers, e1, e2, b1, b2), Nat.add(Nat.mul(outer, inner), k)) == Nat.add(batch_base(inner, outers, e1, e2, b1, b2), _) : Nat}
466      %Equal.sym(Nat, Nat.add(batch_base(inner, outers, e1, e2, b1, b2), 1n+Nat.add(Nat.mul(outer, inner), k)),
467        1n+Nat.add(batch_base(inner, outers, e1, e2, b1, b2), Nat.add(Nat.mul(outer, inner), k)),
468        Arith.add_succ(batch_base(inner, outers, e1, e2, b1, b2), Nat.add(Nat.mul(outer, inner), k))) :
469        {1n+Nat.add(batch_base(inner, outers, e1, e2, b1, b2), Nat.add(Nat.mul(outer, inner), k)) == _ : Nat}
470      {==}
471    case Matrix.Swapped{}:
472      %Index.add_assoc(outers, Nat.mul(k, outers), outer) :
473        {Nat.add(outers, Nat.add(batch_base(inner, outers, e1, e2, b1, b2), Nat.add(Nat.mul(k, outers), outer))) == Nat.add(batch_base(inner, outers, e1, e2, b1, b2), _) : Nat}
474      Nats.add_swap(outers, batch_base(inner, outers, e1, e2, b1, b2), Nat.add(Nat.mul(k, outers), outer))
475
476# A run whose reads start at a cell's first terms reads the cell's logical
477# factor pairs.
478def strided_logical(+fuel: Nat, +dl: Nat, +dr: Nat, -A: Array<F32>, -B: Array<F32>, +lo: Matrix.Order, +ro: Matrix.Order, +inner: Nat,
479  +rows: Nat, +columns: Nat, +l1: Nat, +l2: Nat, +r1: Nat, +r2: Nat, +b1: Nat, +b2: Nat, +row: Nat, +column: Nat, +k: Nat,
480  +lp: Nat, +rp: Nat, values: List<&2, DotFold.Factors>,
481  -el: {lp == Laws.operand_offset(lo, inner, rows, l1, l2, b1, b2, row, k) : Nat},
482  -er: {rp == Laws.operand_offset(ro, inner, columns, r1, r2, b1, b2, column, k) : Nat},
483  -strided: Laws.Strided(fuel, dl, dr, A, B, Matrix.stride(lo, rows), Matrix.stride(ro, columns), lp, rp, values)) ->
484  Parallel.Fact<Laws.Logical(fuel, dl, dr, A, B, lo, ro, inner, rows, columns, l1, l2, r1, r2, b1, b2, row, column, k, values)>:
485  match fuel values:
486    case 0n Nil{}:
487      Parallel.Fact{Unit{}}
488    case 1n+tail DotFold.Factors{lv, rv} <> rest:
489      Parallel.Fact{{\{Laws.StridedStep: left_observed => right_observed => remaining =>
490        Laws.LogicalStep{
491          Equal.trans(Maybe<&2, F32>, Model.observe(F32, dl, A, Refinement.route(dl, U32.from_nat(Laws.operand_offset(lo, inner, rows, l1, l2, b1, b2, row, k)))),
492            Model.observe(F32, dl, A, Refinement.route(dl, U32.from_nat(lp))), Some{lv},
493            Equal.cong(Nat, Maybe<&2, F32>, x => Model.observe(F32, dl, A, Refinement.route(dl, U32.from_nat(x))),
494              Laws.operand_offset(lo, inner, rows, l1, l2, b1, b2, row, k), lp, Equal.sym(Nat, lp, Laws.operand_offset(lo, inner, rows, l1, l2, b1, b2, row, k), el)),
495            left_observed),
496          Equal.trans(Maybe<&2, F32>, Model.observe(F32, dr, B, Refinement.route(dr, U32.from_nat(Laws.operand_offset(ro, inner, columns, r1, r2, b1, b2, column, k)))),
497            Model.observe(F32, dr, B, Refinement.route(dr, U32.from_nat(rp))), Some{rv},
498            Equal.cong(Nat, Maybe<&2, F32>, x => Model.observe(F32, dr, B, Refinement.route(dr, U32.from_nat(x))),
499              Laws.operand_offset(ro, inner, columns, r1, r2, b1, b2, column, k), rp, Equal.sym(Nat, rp, Laws.operand_offset(ro, inner, columns, r1, r2, b1, b2, column, k), er)),
500            right_observed),
501          {\{Parallel.Fact: p => p} : Parallel.Fact<Laws.Logical(tail, dl, dr, A, B, lo, ro, inner, rows, columns, l1, l2, r1, r2, b1, b2, row, column, 1n+k, rest)> ->
502            Laws.Logical(tail, dl, dr, A, B, lo, ro, inner, rows, columns, l1, l2, r1, r2, b1, b2, row, column, 1n+k, rest)}(
503            strided_logical(tail, dl, dr, A, B, lo, ro, inner, rows, columns, l1, l2, r1, r2, b1, b2, row, column, 1n+k,
504              Nat.add(Matrix.stride(lo, rows), lp), Nat.add(Matrix.stride(ro, columns), rp), rest,
505              Equal.trans(Nat, Nat.add(Matrix.stride(lo, rows), lp), Nat.add(Matrix.stride(lo, rows), Laws.operand_offset(lo, inner, rows, l1, l2, b1, b2, row, k)),
506                Laws.operand_offset(lo, inner, rows, l1, l2, b1, b2, row, 1n+k),
507                Equal.cong(Nat, Nat, x => Nat.add(Matrix.stride(lo, rows), x), lp, Laws.operand_offset(lo, inner, rows, l1, l2, b1, b2, row, k), el),
508                operand_step(lo, inner, rows, l1, l2, b1, b2, row, k)),
509              Equal.trans(Nat, Nat.add(Matrix.stride(ro, columns), rp), Nat.add(Matrix.stride(ro, columns), Laws.operand_offset(ro, inner, columns, r1, r2, b1, b2, column, k)),
510                Laws.operand_offset(ro, inner, columns, r1, r2, b1, b2, column, 1n+k),
511                Equal.cong(Nat, Nat, x => Nat.add(Matrix.stride(ro, columns), x), rp, Laws.operand_offset(ro, inner, columns, r1, r2, b1, b2, column, k), er),
512                operand_step(ro, inner, columns, r1, r2, b1, b2, column, k)),
513              remaining))}} :
514        Laws.StridedStep<tail, dl, dr, A, B, Matrix.stride(lo, rows), Matrix.stride(ro, columns), lp, rp, lv, rv, rest> ->
515        Laws.LogicalStep<tail, dl, dr, A, B, lo, ro, inner, rows, columns, l1, l2, r1, r2, b1, b2, row, column, k, lv, rv, rest>}(strided)}
516    case 0n DotFold.Factors{lv, rv} <> rest:
517      Parallel.Fact{Empty.absurd(Laws.Logical(0n, dl, dr, A, B, lo, ro, inner, rows, columns, l1, l2, r1, r2, b1, b2, row, column, k, DotFold.Factors{lv, rv} <> rest), strided)}
518    case 1n+tail Nil{}:
519      Parallel.Fact{Empty.absurd(Laws.Logical(1n+tail, dl, dr, A, B, lo, ro, inner, rows, columns, l1, l2, r1, r2, b1, b2, row, column, k, Nil{}), strided)}
520
521def CellIs(+dl: Nat, +dr: Nat, -A: Array<F32>, -B: Array<F32>, +lo: Matrix.Order, +ro: Matrix.Order, +inner: Nat, +rows: Nat,
522  +columns: Nat, +l1: Nat, +l2: Nat, +r1: Nat, +r2: Nat, +o2: Nat, +count: Nat, +position: Nat, -value: Maybe<&2, F32>) -> Type:
523  Laws.CellSpec(dl, dr, A, B, lo, ro, inner, rows, columns, l1, l2, r1, r2, o2, position, Nat.is_lt(position, count), value)
524
525def cell_shift(+dl: Nat, +dr: Nat, -A: Array<F32>, -B: Array<F32>, +lo: Matrix.Order, +ro: Matrix.Order, +inner: Nat, +rows: Nat,
526  +columns: Nat, +l1: Nat, +l2: Nat, +r1: Nat, +r2: Nat, +o2: Nat, +count: Nat, -position: Nat, -other: Nat, -value: Maybe<&2, F32>,
527  -same: {position == other : Nat}, -at: CellIs(dl, dr, A, B, lo, ro, inner, rows, columns, l1, l2, r1, r2, o2, count, position, value)) ->
528  Parallel.Fact<CellIs(dl, dr, A, B, lo, ro, inner, rows, columns, l1, l2, r1, r2, o2, count, other, value)>:
529  Parallel.Fact{
530    %same : CellIs(dl, dr, A, B, lo, ro, inner, rows, columns, l1, l2, r1, r2, o2, count, _, value)
531    at}
532
533def cell_below(+dl: Nat, +dr: Nat, -A: Array<F32>, -B: Array<F32>, +lo: Matrix.Order, +ro: Matrix.Order, +inner: Nat, +rows: Nat,
534  +columns: Nat, +l1: Nat, +l2: Nat, +r1: Nat, +r2: Nat, +o2: Nat, +count: Nat, -position: Nat, -value: Maybe<&2, F32>,
535  -below: {Nat.is_lt(position, count) == True{} : Bool},
536  -at: CellIs(dl, dr, A, B, lo, ro, inner, rows, columns, l1, l2, r1, r2, o2, count, position, value)) ->
537  Parallel.Fact<Laws.CellAt<dl, dr, A, B, lo, ro, inner, rows, columns, l1, l2, r1, r2, o2, position, value>>:
538  Parallel.Fact{
539    %below : Laws.CellSpec(dl, dr, A, B, lo, ro, inner, rows, columns, l1, l2, r1, r2, o2, position, _, value)
540    at}
541
542def Laws.batched_build_logical_cells(depth, splits, dl, dr, A, B, lo, ro, inner, rows, columns, l1, l2, r1, r2, o2, count, base, la, ra, el, er,
543  left_valid, right_valid, path, b1, b2, row, column, row_below, column_below, batch_below, at, inside):
544  -cells = {\{Laws.Holds: q => q} :
545    Laws.Holds<Laws.Cells(depth, dl, dr, A, B, lo, ro, inner, rows, columns, l1, l2, r1, r2, o2, count, base,
546      Batched.build(depth, splits, Laws.layout_of(lo, ro, inner, rows, columns, l1, l2, r1, r2, o2), count, base, la, ra))> ->
547    Laws.Cells(depth, dl, dr, A, B, lo, ro, inner, rows, columns, l1, l2, r1, r2, o2, count, base,
548      Batched.build(depth, splits, Laws.layout_of(lo, ro, inner, rows, columns, l1, l2, r1, r2, o2), count, base, la, ra))}(
549    Laws.batched_build_cells(depth, splits, dl, dr, A, B, lo, ro, inner, rows, columns, l1, l2, r1, r2, o2, count, base, la, ra, el, er,
550      left_valid, right_valid))
551  -shifted = {\{Parallel.Fact: q => q} :
552    Parallel.Fact<CellIs(dl, dr, A, B, lo, ro, inner, rows, columns, l1, l2, r1, r2, o2, count, Laws.cell_position(o2, rows, columns, b1, b2, row, column),
553      Model.observe(F32, depth, Batched.build(depth, splits, Laws.layout_of(lo, ro, inner, rows, columns, l1, l2, r1, r2, o2), count, base, la, ra), path))> ->
554    CellIs(dl, dr, A, B, lo, ro, inner, rows, columns, l1, l2, r1, r2, o2, count, Laws.cell_position(o2, rows, columns, b1, b2, row, column),
555      Model.observe(F32, depth, Batched.build(depth, splits, Laws.layout_of(lo, ro, inner, rows, columns, l1, l2, r1, r2, o2), count, base, la, ra), path))}(
556    cell_shift(dl, dr, A, B, lo, ro, inner, rows, columns, l1, l2, r1, r2, o2, count, Nat.add(Model.offset(depth, path), base),
557      Laws.cell_position(o2, rows, columns, b1, b2, row, column),
558      Model.observe(F32, depth, Batched.build(depth, splits, Laws.layout_of(lo, ro, inner, rows, columns, l1, l2, r1, r2, o2), count, base, la, ra), path),
559      at, cells(path)))
560  -cell = {\{Parallel.Fact: q => q} :
561    Parallel.Fact<Laws.CellAt<dl, dr, A, B, lo, ro, inner, rows, columns, l1, l2, r1, r2, o2, Laws.cell_position(o2, rows, columns, b1, b2, row, column),
562      Model.observe(F32, depth, Batched.build(depth, splits, Laws.layout_of(lo, ro, inner, rows, columns, l1, l2, r1, r2, o2), count, base, la, ra), path)>> ->
563    Laws.CellAt<dl, dr, A, B, lo, ro, inner, rows, columns, l1, l2, r1, r2, o2, Laws.cell_position(o2, rows, columns, b1, b2, row, column),
564      Model.observe(F32, depth, Batched.build(depth, splits, Laws.layout_of(lo, ro, inner, rows, columns, l1, l2, r1, r2, o2), count, base, la, ra), path)>}(
565    cell_below(dl, dr, A, B, lo, ro, inner, rows, columns, l1, l2, r1, r2, o2, count, Laws.cell_position(o2, rows, columns, b1, b2, row, column),
566      Model.observe(F32, depth, Batched.build(depth, splits, Laws.layout_of(lo, ro, inner, rows, columns, l1, l2, r1, r2, o2), count, base, la, ra), path),
567      inside, shifted))
568  Laws.Holds{{\{Laws.CellAt: word => observed => values => logical => fold =>
569    Laws.LogicalCell{word, observed, values,
570      {\{Parallel.Fact: q => q} :
571        Parallel.Fact<Laws.Logical(inner, dl, dr, A, B, lo, ro, inner, rows, columns, l1, l2, r1, r2, b1, b2, row, column, 0n, values)> ->
572        Laws.Logical(inner, dl, dr, A, B, lo, ro, inner, rows, columns, l1, l2, r1, r2, b1, b2, row, column, 0n, values)}(
573        strided_logical(inner, dl, dr, A, B, lo, ro, inner, rows, columns, l1, l2, r1, r2, b1, b2, row, column, 0n,
574          Batched.left_start(Laws.layout_of(lo, ro, inner, rows, columns, l1, l2, r1, r2, o2), Laws.cell_position(o2, rows, columns, b1, b2, row, column)),
575          Batched.right_start(Laws.layout_of(lo, ro, inner, rows, columns, l1, l2, r1, r2, o2), Laws.cell_position(o2, rows, columns, b1, b2, row, column)),
576          values,
577          left_start_spec(lo, ro, inner, rows, columns, l1, l2, r1, r2, o2, b1, b2, row, column, row_below, column_below, batch_below),
578          right_start_spec(lo, ro, inner, rows, columns, l1, l2, r1, r2, o2, b1, b2, row, column, row_below, column_below, batch_below),
579          logical)),
580      fold}} :
581    Laws.CellAt<dl, dr, A, B, lo, ro, inner, rows, columns, l1, l2, r1, r2, o2, Laws.cell_position(o2, rows, columns, b1, b2, row, column),
582      Model.observe(F32, depth, Batched.build(depth, splits, Laws.layout_of(lo, ro, inner, rows, columns, l1, l2, r1, r2, o2), count, base, la, ra), path)> ->
583    Laws.LogicalCell<dl, dr, A, B, lo, ro, inner, rows, columns, l1, l2, r1, r2, b1, b2, row, column,
584      Model.observe(F32, depth, Batched.build(depth, splits, Laws.layout_of(lo, ro, inner, rows, columns, l1, l2, r1, r2, o2), count, base, la, ra), path)>}(cell)}
585
586# The plan's geometry.
587
588def AxisMeaning(+l: U32, +r: U32, result: Result<&2, &2, Matrix.Error, U32>) -> Data:
589  match result:
590    case Fail{error}:
591      Unit
592    case Done{o}:
593      {U32.to_nat(o) == Laws.extent(U32.to_nat(l), U32.to_nat(r)) : Nat}
594
595def same_extent(+a: Nat, +b: Nat, +same: {a == b : Nat}, decision: Shape.BoolCase<Nat.is_eq(a, 1n)>) ->
596  {a == Laws.extent(a, b) : Nat}:
597  match decision:
598    case Shape.BTrue{equal}:
599      %equal : {a == Bool.pick(Nat, _, b, a) : Nat}
600      same
601    case Shape.BFalse{equal}:
602      %equal : {a == Bool.pick(Nat, _, b, a) : Nat}
603      {==}
604
605def axis_right(+l: U32, +r: U32, +left_one: {False{} == Nat.is_eq(U32.to_nat(l), 1n) : Bool},
606  decision: Shape.BoolCase<Nat.is_eq(U32.to_nat(r), 1n)>) ->
607  AxisMeaning(l, r, Batched.stretched(l, r, False{}, False{}, Nat.is_eq(U32.to_nat(r), 1n))):
608  match decision:
609    case Shape.BTrue{equal}:
610      %equal : AxisMeaning(l, r, Batched.stretched(l, r, False{}, False{}, _))
611      %left_one : {U32.to_nat(l) == Bool.pick(Nat, _, U32.to_nat(r), U32.to_nat(l)) : Nat}
612      {==}
613    case Shape.BFalse{equal}:
614      %equal : AxisMeaning(l, r, Batched.stretched(l, r, False{}, False{}, _))
615      Unit{}
616
617def axis_left(+l: U32, +r: U32, decision: Shape.BoolCase<Nat.is_eq(U32.to_nat(l), 1n)>) ->
618  AxisMeaning(l, r, Batched.stretched(l, r, False{}, Nat.is_eq(U32.to_nat(l), 1n), Nat.is_eq(U32.to_nat(r), 1n))):
619  match decision:
620    case Shape.BTrue{+equal}:
621      %equal : AxisMeaning(l, r, Batched.stretched(l, r, False{}, _, Nat.is_eq(U32.to_nat(r), 1n)))
622      %equal : {U32.to_nat(r) == Bool.pick(Nat, _, U32.to_nat(r), U32.to_nat(l)) : Nat}
623      {==}
624    case Shape.BFalse{+equal}:
625      %equal : AxisMeaning(l, r, Batched.stretched(l, r, False{}, _, Nat.is_eq(U32.to_nat(r), 1n)))
626      axis_right(l, r, equal, Shape.bool_case(Nat.is_eq(U32.to_nat(r), 1n)))
627
628def axis_same(+l: U32, +r: U32, decision: Shape.BoolCase<Nat.is_eq(U32.to_nat(l), U32.to_nat(r))>) ->
629  AxisMeaning(l, r, Batched.stretched(l, r, Nat.is_eq(U32.to_nat(l), U32.to_nat(r)), Nat.is_eq(U32.to_nat(l), 1n), Nat.is_eq(U32.to_nat(r), 1n))):
630  match decision:
631    case Shape.BTrue{+equal}:
632      %equal : AxisMeaning(l, r, Batched.stretched(l, r, _, Nat.is_eq(U32.to_nat(l), 1n), Nat.is_eq(U32.to_nat(r), 1n)))
633      same_extent(U32.to_nat(l), U32.to_nat(r), Lifecycle.natural_equal(U32.to_nat(l), U32.to_nat(r), equal),
634        Shape.bool_case(Nat.is_eq(U32.to_nat(l), 1n)))
635    case Shape.BFalse{equal}:
636      %equal : AxisMeaning(l, r, Batched.stretched(l, r, _, Nat.is_eq(U32.to_nat(l), 1n), Nat.is_eq(U32.to_nat(r), 1n)))
637      axis_left(l, r, Shape.bool_case(Nat.is_eq(U32.to_nat(l), 1n)))
638
639# One batch axis of the plan has the extent PyTorch's broadcast gives it.
640def axis_meaning(+l: U32, +r: U32) -> AxisMeaning(l, r, Batched.axis(l, r)):
641  axis_same(l, r, Shape.bool_case(Nat.is_eq(U32.to_nat(l), U32.to_nat(r))))
642
643# The plan's batch extents as naturals, as `Batched.batch_dims` lists them.
644def batch_nats(+rank: Nat, +first: Nat, +second: Nat) -> List<&2, Nat>:
645  match rank:
646    case 0n:
647      []
648    case 1n+rest:
649      match rest:
650        case 0n:
651          [second]
652        case 1n+more:
653          [first, second]
654
655def naturals_append(+xs: List<&2, U32>, +ys: List<&2, U32>) ->
656  {Shape.naturals(List.append(&2, U32, xs, ys)) == List.append(&2, Nat, Shape.naturals(xs), Shape.naturals(ys)) : List<&2, Nat>}:
657  match xs:
658    case Nil{}:
659      {==}
660    case +head <> tail:
661      Equal.cong(List<&2, Nat>, List<&2, Nat>, rest => U32.to_nat(head) <> rest,
662        Shape.naturals(List.append(&2, U32, tail, ys)), List.append(&2, Nat, Shape.naturals(tail), Shape.naturals(ys)), naturals_append(tail, ys))
663
664def batch_part(+rank: Nat, +o1: U32, +o2: U32, +first: Nat, +second: Nat,
665  +e1: {U32.to_nat(o1) == first : Nat}, +e2: {U32.to_nat(o2) == second : Nat}) ->
666  {Shape.naturals(Batched.batch_dims(rank, o1, o2)) == batch_nats(rank, first, second) : List<&2, Nat>}:
667  match rank:
668    case 0n:
669      {==}
670    case 1n+rest:
671      match rest:
672        case 0n:
673          %Equal.sym(Nat, U32.to_nat(o2), second, e2) : {[_] == [second] : List<&2, Nat>}
674          {==}
675        case 1n+more:
676          %Equal.sym(Nat, U32.to_nat(o1), first, e1) : {[_, U32.to_nat(o2)] == [first, second] : List<&2, Nat>}
677          %Equal.sym(Nat, U32.to_nat(o2), second, e2) : {[first, _] == [first, second] : List<&2, Nat>}
678          {==}
679
680def matrix_part(+rows: U32, +lv: Bool, +columns: U32, +rv: Bool) ->
681  {Shape.naturals(Batched.matrix_dims(rows, lv, columns, rv)) ==
682    List.append(&2, Nat, Laws.kept(lv, U32.to_nat(rows)), Laws.kept(rv, U32.to_nat(columns))) : List<&2, Nat>}:
683  match lv rv:
684    case True{} True{}:
685      {==}
686    case True{} False{}:
687      {==}
688    case False{} True{}:
689      {==}
690    case False{} False{}:
691      {==}
692
693def output_nats(+rank: Nat, +o1: U32, +o2: U32, +first: Nat, +second: Nat,
694  +e1: {U32.to_nat(o1) == first : Nat}, +e2: {U32.to_nat(o2) == second : Nat}, +rows: U32, +lv: Bool, +columns: U32, +rv: Bool) ->
695  {Shape.naturals(List.append(&2, U32, Batched.batch_dims(rank, o1, o2), Batched.matrix_dims(rows, lv, columns, rv))) ==
696    List.append(&2, Nat, batch_nats(rank, first, second),
697      List.append(&2, Nat, Laws.kept(lv, U32.to_nat(rows)), Laws.kept(rv, U32.to_nat(columns)))) : List<&2, Nat>}:
698  %Equal.sym(List<&2, Nat>, Shape.naturals(List.append(&2, U32, Batched.batch_dims(rank, o1, o2), Batched.matrix_dims(rows, lv, columns, rv))),
699    List.append(&2, Nat, Shape.naturals(Batched.batch_dims(rank, o1, o2)), Shape.naturals(Batched.matrix_dims(rows, lv, columns, rv))),
700    naturals_append(Batched.batch_dims(rank, o1, o2), Batched.matrix_dims(rows, lv, columns, rv))) :
701    {_ == List.append(&2, Nat, batch_nats(rank, first, second),
702      List.append(&2, Nat, Laws.kept(lv, U32.to_nat(rows)), Laws.kept(rv, U32.to_nat(columns)))) : List<&2, Nat>}
703  %Equal.sym(List<&2, Nat>, Shape.naturals(Batched.batch_dims(rank, o1, o2)), batch_nats(rank, first, second), batch_part(rank, o1, o2, first, second, e1, e2)) :
704    {List.append(&2, Nat, _, Shape.naturals(Batched.matrix_dims(rows, lv, columns, rv))) == List.append(&2, Nat, batch_nats(rank, first, second),
705      List.append(&2, Nat, Laws.kept(lv, U32.to_nat(rows)), Laws.kept(rv, U32.to_nat(columns)))) : List<&2, Nat>}
706  %Equal.sym(List<&2, Nat>, Shape.naturals(Batched.matrix_dims(rows, lv, columns, rv)),
707    List.append(&2, Nat, Laws.kept(lv, U32.to_nat(rows)), Laws.kept(rv, U32.to_nat(columns))), matrix_part(rows, lv, columns, rv)) :
708    {List.append(&2, Nat, batch_nats(rank, first, second), _) == List.append(&2, Nat, batch_nats(rank, first, second),
709      List.append(&2, Nat, Laws.kept(lv, U32.to_nat(rows)), Laws.kept(rv, U32.to_nat(columns)))) : List<&2, Nat>}
710  {==}
711
712def pair_nats(+a: Nat, +b: Nat, +c: Nat, +d: Nat, +ac: {a == c : Nat}, +bd: {b == d : Nat}) -> {[a, b] == [c, d] : List<&2, Nat>}:
713  %Equal.sym(Nat, a, c, ac) : {[_, b] == [c, d] : List<&2, Nat>}
714  %Equal.sym(Nat, b, d, bd) : {[c, _] == [c, d] : List<&2, Nat>}
715  {==}
716
717# What the plan of two operands of these fields holds, in the operands'
718# terms.
719def ExpectedOf(+lr: Nat, +l1: U32, +l2: U32, +rows: U32, +inner: U32, +lv: Bool, +rr: Nat, +r1: U32, +r2: U32, +columns: U32,
720  +rinner: U32, +rv: Bool, result: Result<&2, &2, Matrix.Error, Batched.Plan>) -> Data:
721  Laws.GeometryOf(
722    List.append(&2, Nat, batch_nats(Batched.larger(lr, rr), Laws.extent(U32.to_nat(l1), U32.to_nat(r1)), Laws.extent(U32.to_nat(l2), U32.to_nat(r2))),
723      List.append(&2, Nat, Laws.kept(lv, U32.to_nat(rows)), Laws.kept(rv, U32.to_nat(columns)))),
724    [U32.to_nat(l1), U32.to_nat(l2), U32.to_nat(rows), U32.to_nat(inner)],
725    [U32.to_nat(r1), U32.to_nat(r2), U32.to_nat(rinner), U32.to_nat(columns)],
726    [Laws.extent(U32.to_nat(l1), U32.to_nat(r1)), Laws.extent(U32.to_nat(l2), U32.to_nat(r2))],
727    (lv, rv), result)
728
729def geometry_shaped(+lr: Nat, +l1: U32, +l2: U32, +rows: U32, +inner: U32, +lv: Bool, +rr: Nat, +r1: U32, +r2: U32, +columns: U32,
730  +rinner: U32, +rv: Bool, +o1: U32, +o2: U32, +same: {U32.to_nat(inner) == U32.to_nat(rinner) : Nat},
731  +e1: {U32.to_nat(o1) == Laws.extent(U32.to_nat(l1), U32.to_nat(r1)) : Nat},
732  +e2: {U32.to_nat(o2) == Laws.extent(U32.to_nat(l2), U32.to_nat(r2)) : Nat},
733  result: Shape.ConstructResult<List.append(&2, U32, Batched.batch_dims(Batched.larger(lr, rr), o1, o2), Batched.matrix_dims(rows, lv, columns, rv))>) ->
734  ExpectedOf(lr, l1, l2, rows, inner, lv, rr, r1, r2, columns, rinner, rv,
735    Batched.planned_shape(Batched.batch_dims(Batched.larger(lr, rr), o1, o2), rows, inner, columns, l1, l2, r1, r2, o1, o2, lv, rv,
736      List.append(&2, U32, Batched.batch_dims(Batched.larger(lr, rr), o1, o2), Batched.matrix_dims(rows, lv, columns, rv)), result)):
737  match result:
738    case Shape.SError{error}:
739      Unit{}
740    case Shape.SBuilt{+shape, +preserved}:
741      Laws.Geometry{
742        Equal.trans(List<&2, Nat>, Shape.naturals(Shape.dimensions(shape)),
743          Shape.naturals(List.append(&2, U32, Batched.batch_dims(Batched.larger(lr, rr), o1, o2), Batched.matrix_dims(rows, lv, columns, rv))),
744          List.append(&2, Nat, batch_nats(Batched.larger(lr, rr), Laws.extent(U32.to_nat(l1), U32.to_nat(r1)), Laws.extent(U32.to_nat(l2), U32.to_nat(r2))),
745            List.append(&2, Nat, Laws.kept(lv, U32.to_nat(rows)), Laws.kept(rv, U32.to_nat(columns)))),
746          Equal.cong(List<&2, U32>, List<&2, Nat>, dims => Shape.naturals(dims), Shape.dimensions(shape),
747            List.append(&2, U32, Batched.batch_dims(Batched.larger(lr, rr), o1, o2), Batched.matrix_dims(rows, lv, columns, rv)), preserved),
748          output_nats(Batched.larger(lr, rr), o1, o2, Laws.extent(U32.to_nat(l1), U32.to_nat(r1)), Laws.extent(U32.to_nat(l2), U32.to_nat(r2)),
749            e1, e2, rows, lv, columns, rv)),
750        {==},
751        Equal.cong(Nat, List<&2, Nat>, x => [U32.to_nat(r1), U32.to_nat(r2), x, U32.to_nat(columns)], U32.to_nat(inner), U32.to_nat(rinner), same),
752        pair_nats(U32.to_nat(o1), U32.to_nat(o2), Laws.extent(U32.to_nat(l1), U32.to_nat(r1)), Laws.extent(U32.to_nat(l2), U32.to_nat(r2)), e1, e2),
753        {==}}
754
755def geometry_batched(+lr: Nat, +l1: U32, +l2: U32, +rows: U32, +inner: U32, +lv: Bool, +rr: Nat, +r1: U32, +r2: U32, +columns: U32,
756  +rinner: U32, +rv: Bool, +same: {U32.to_nat(inner) == U32.to_nat(rinner) : Nat},
757  first: Result<&2, &2, Matrix.Error, U32>, second: Result<&2, &2, Matrix.Error, U32>,
758  +m1: AxisMeaning(l1, r1, first), +m2: AxisMeaning(l2, r2, second)) ->
759  ExpectedOf(lr, l1, l2, rows, inner, lv, rr, r1, r2, columns, rinner, rv,
760    Batched.batched(Batched.larger(lr, rr), rows, inner, columns, l1, l2, r1, r2, lv, rv, first, second)):
761  match first second:
762    case Fail{error} second:
763      Unit{}
764    case Done{o1} Fail{error}:
765      Unit{}
766    case Done{+o1} Done{+o2}:
767      geometry_shaped(lr, l1, l2, rows, inner, lv, rr, r1, r2, columns, rinner, rv, o1, o2, same, m1, m2,
768        Shape.construct(List.append(&2, U32, Batched.batch_dims(Batched.larger(lr, rr), o1, o2), Batched.matrix_dims(rows, lv, columns, rv))))
769
770def geometry_decided(+lr: Nat, +l1: U32, +l2: U32, +rows: U32, +inner: U32, +lv: Bool, +rr: Nat, +r1: U32, +r2: U32, +columns: U32,
771  +rinner: U32, +rv: Bool, decision: Shape.BoolCase<Nat.is_eq(U32.to_nat(inner), U32.to_nat(rinner))>) ->
772  ExpectedOf(lr, l1, l2, rows, inner, lv, rr, r1, r2, columns, rinner, rv,
773    Batched.compatible(Batched.Operand{Batched.Batch{lr, l1, l2}, rows, inner, lv}, Batched.Operand{Batched.Batch{rr, r1, r2}, columns, rinner, rv},
774      Nat.is_eq(U32.to_nat(inner), U32.to_nat(rinner)))):
775  match decision:
776    case Shape.BFalse{equal}:
777      %equal : ExpectedOf(lr, l1, l2, rows, inner, lv, rr, r1, r2, columns, rinner, rv,
778        Batched.compatible(Batched.Operand{Batched.Batch{lr, l1, l2}, rows, inner, lv}, Batched.Operand{Batched.Batch{rr, r1, r2}, columns, rinner, rv}, _))
779      Unit{}
780    case Shape.BTrue{+equal}:
781      %equal : ExpectedOf(lr, l1, l2, rows, inner, lv, rr, r1, r2, columns, rinner, rv,
782        Batched.compatible(Batched.Operand{Batched.Batch{lr, l1, l2}, rows, inner, lv}, Batched.Operand{Batched.Batch{rr, r1, r2}, columns, rinner, rv}, _))
783      geometry_batched(lr, l1, l2, rows, inner, lv, rr, r1, r2, columns, rinner, rv,
784        Lifecycle.natural_equal(U32.to_nat(inner), U32.to_nat(rinner), equal),
785        Batched.axis(l1, r1), Batched.axis(l2, r2), axis_meaning(l1, r1), axis_meaning(l2, r2))
786
787def geometry_ops(+lr: Nat, +l1: U32, +l2: U32, +rows: U32, +inner: U32, +lv: Bool, +rr: Nat, +r1: U32, +r2: U32, +columns: U32,
788  +rinner: U32, +rv: Bool) ->
789  ExpectedOf(lr, l1, l2, rows, inner, lv, rr, r1, r2, columns, rinner, rv,
790    Batched.compatible(Batched.Operand{Batched.Batch{lr, l1, l2}, rows, inner, lv}, Batched.Operand{Batched.Batch{rr, r1, r2}, columns, rinner, rv},
791      Nat.is_eq(U32.to_nat(inner), U32.to_nat(rinner)))):
792  geometry_decided(lr, l1, l2, rows, inner, lv, rr, r1, r2, columns, rinner, rv, Shape.bool_case(Nat.is_eq(U32.to_nat(inner), U32.to_nat(rinner))))
793
794# The plan of each pair of ranks 1 to 4 is that of its two operands; any
795# other rank is refused.
796def Laws.batched_plan_geometry(left, right):
797  match left:
798    case Nil{}:
799      Unit{}
800    case +k <> Nil{}:
801      match right:
802        case Nil{}:
803          Unit{}
804        case +q <> Nil{}:
805          geometry_ops(0n, 1, 1, 1, k, True{}, 0n, 1, 1, 1, q, True{})
806        case +q <> +n <> Nil{}:
807          geometry_ops(0n, 1, 1, 1, k, True{}, 0n, 1, 1, n, q, False{})
808        case +c <> +q <> +n <> Nil{}:
809          geometry_ops(0n, 1, 1, 1, k, True{}, 1n, 1, c, n, q, False{})
810        case +c <> +d <> +q <> +n <> Nil{}:
811          geometry_ops(0n, 1, 1, 1, k, True{}, 2n, c, d, n, q, False{})
812        case +w <> +x <> +y <> +z <> +t <> rest:
813          Unit{}
814    case +m <> +k <> Nil{}:
815      match right:
816        case Nil{}:
817          Unit{}
818        case +q <> Nil{}:
819          geometry_ops(0n, 1, 1, m, k, False{}, 0n, 1, 1, 1, q, True{})
820        case +q <> +n <> Nil{}:
821          geometry_ops(0n, 1, 1, m, k, False{}, 0n, 1, 1, n, q, False{})
822        case +c <> +q <> +n <> Nil{}:
823          geometry_ops(0n, 1, 1, m, k, False{}, 1n, 1, c, n, q, False{})
824        case +c <> +d <> +q <> +n <> Nil{}:
825          geometry_ops(0n, 1, 1, m, k, False{}, 2n, c, d, n, q, False{})
826        case +w <> +x <> +y <> +z <> +t <> rest:
827          Unit{}
828    case +b <> +m <> +k <> Nil{}:
829      match right:
830        case Nil{}:
831          Unit{}
832        case +q <> Nil{}:
833          geometry_ops(1n, 1, b, m, k, False{}, 0n, 1, 1, 1, q, True{})
834        case +q <> +n <> Nil{}:
835          geometry_ops(1n, 1, b, m, k, False{}, 0n, 1, 1, n, q, False{})
836        case +c <> +q <> +n <> Nil{}:
837          geometry_ops(1n, 1, b, m, k, False{}, 1n, 1, c, n, q, False{})
838        case +c <> +d <> +q <> +n <> Nil{}:
839          geometry_ops(1n, 1, b, m, k, False{}, 2n, c, d, n, q, False{})
840        case +w <> +x <> +y <> +z <> +t <> rest:
841          Unit{}
842    case +a <> +b <> +m <> +k <> Nil{}:
843      match right:
844        case Nil{}:
845          Unit{}
846        case +q <> Nil{}:
847          geometry_ops(2n, a, b, m, k, False{}, 0n, 1, 1, 1, q, True{})
848        case +q <> +n <> Nil{}:
849          geometry_ops(2n, a, b, m, k, False{}, 0n, 1, 1, n, q, False{})
850        case +c <> +q <> +n <> Nil{}:
851          geometry_ops(2n, a, b, m, k, False{}, 1n, 1, c, n, q, False{})
852        case +c <> +d <> +q <> +n <> Nil{}:
853          geometry_ops(2n, a, b, m, k, False{}, 2n, c, d, n, q, False{})
854        case +w <> +x <> +y <> +z <> +t <> rest:
855          Unit{}
856    case +a <> +b <> +c <> +d <> +e <> rest:
857      Unit{}