1	import Base
2	import ../shape/core.bend as Shape
3	import ../storage/model.bend as Model
4	import ../storage/refinement.bend as Refinement
5	import ./core.bend as Core
6	import ./matmul.bend as Matrix
7	import ./batched_matmul.bend as Batched
8	import ./dot_fold.bend as DotFold
9	
10	# The laws of `Batched.matmul`, PyTorch's torch.matmul: which call answers
11	# a pair of operands, which pairs it refuses, and what every cell of its
12	# parallel build holds.
13	
14	type Holds<-P: Type> is Data:
15	  Holds{-proof: P}
16	
17	# Two rank-2 operands are `Matrix.matmul` itself: the same result, words,
18	# errors and returned inputs, so every law of `matmul` holds for them and
19	# their words do not change.
20	law batched_matmul_matrices:
21	  for +left: Shape.Shape
22	  for +right: Shape.Shape
23	  for a: Core.Tensor<left>
24	  for b: Core.Tensor<right>
25	  for +rows: U32
26	  for +inner: U32
27	  for +right_inner: U32
28	  for +columns: U32
29	  for left_matrix: {Shape.dimensions(left) == [rows, inner] : List<&2, U32>}
30	  for right_matrix: {Shape.dimensions(right) == [right_inner, columns] : List<&2, U32>}
31	  {Batched.matmul(Core.Packed{left, a}, Core.Packed{right, b}) == Matrix.matmul(Core.Packed{left, a}, Core.Packed{right, b}) :
32	    Result<&2, &1, Matrix.Error, Core.Dynamic & Core.Dynamic & Core.Dynamic>}
33	
34	# A rank-0 left operand fails with RankMismatch, whatever the right operand
35	# is, as PyTorch refuses it ("both arguments to matmul need to be at least
36	# 1D").
37	law batched_matmul_left_scalar:
38	  for +left: Shape.Shape
39	  for +right: Shape.Shape
40	  for a: Core.Tensor<left>
41	  for b: Core.Tensor<right>
42	  for scalar: {Shape.dimensions(left) == [] : List<&2, U32>}
43	  {Batched.matmul(Core.Packed{left, a}, Core.Packed{right, b}) == Fail{Matrix.RankMismatch{}} :
44	    Result<&2, &1, Matrix.Error, Core.Dynamic & Core.Dynamic & Core.Dynamic>}
45	
46	# A rank-0 right operand fails with RankMismatch, whatever the left operand
47	# is.
48	law batched_matmul_right_scalar:
49	  for +left: Shape.Shape
50	  for +right: Shape.Shape
51	  for a: Core.Tensor<left>
52	  for b: Core.Tensor<right>
53	  for scalar: {Shape.dimensions(right) == [] : List<&2, U32>}
54	  {Batched.matmul(Core.Packed{left, a}, Core.Packed{right, b}) == Fail{Matrix.RankMismatch{}} :
55	    Result<&2, &1, Matrix.Error, Core.Dynamic & Core.Dynamic & Core.Dynamic>}
56	
57	# A vector [k] times a vector [k'] with k != k' fails with
58	# InnerDimensionMismatch: torch.matmul's dot product refuses unequal
59	# lengths.
60	law batched_matmul_dot_inner_error:
61	  for +left: Shape.Shape
62	  for +right: Shape.Shape
63	  for a: Core.Tensor<left>
64	  for b: Core.Tensor<right>
65	  for +inner: U32
66	  for +right_inner: U32
67	  for left_vector: {Shape.dimensions(left) == [inner] : List<&2, U32>}
68	  for right_vector: {Shape.dimensions(right) == [right_inner] : List<&2, U32>}
69	  for mismatch: {U32.to_nat(inner) != U32.to_nat(right_inner) : Nat}
70	  {Batched.matmul(Core.Packed{left, a}, Core.Packed{right, b}) == Fail{Matrix.InnerDimensionMismatch{}} :
71	    Result<&2, &1, Matrix.Error, Core.Dynamic & Core.Dynamic & Core.Dynamic>}
72	
73	# The cells of a build: the layout's fields, named so a law can take them.
74	def layout_of(+lo: Matrix.Order, +ro: Matrix.Order, +inner: Nat, +rows: Nat, +columns: Nat,
75	  +l1: Nat, +l2: Nat, +r1: Nat, +r2: Nat, +o2: Nat) -> Batched.Layout:
76	  Batched.Layout{lo, ro, inner, rows, columns, l1, l2, r1, r2, o2}
77	
78	# The factor pairs a run of `fuel` terms reads: A at lp, lp + ls,
79	# lp + 2 ls, ... and B at rp, rp + rs, ..., in that order.
80	def Strided(+fuel: Nat, +dl: Nat, +dr: Nat, -A: Array<F32>, -B: Array<F32>, +ls: Nat, +rs: Nat, +lp: Nat, +rp: Nat,
81	  values: List<&2, DotFold.Factors>) -> Data:
82	  match fuel values:
83	    case 0n Nil{}:
84	      Unit
85	    case 1n+tail DotFold.Factors{lv, rv} <> rest:
86	      StridedStep<tail, dl, dr, A, B, ls, rs, lp, rp, lv, rv, rest>
87	    case 0n DotFold.Factors{lv, rv} <> rest:
88	      Empty
89	    case 1n+tail Nil{}:
90	      Empty
91	
92	type StridedStep<-fuel: Nat, -dl: Nat, -dr: Nat, -A: Array<F32>, -B: Array<F32>, -ls: Nat, -rs: Nat, -lp: Nat, -rp: Nat,
93	  -lv: F32, -rv: F32, -values: List<&2, DotFold.Factors>> is Data:
94	  StridedStep{-left_observed: {Model.observe(F32, dl, A, Refinement.route(dl, U32.from_nat(lp))) == Some{lv} : Maybe<&2, F32>},
95	    -right_observed: {Model.observe(F32, dr, B, Refinement.route(dr, U32.from_nat(rp))) == Some{rv} : Maybe<&2, F32>},
96	    -remaining: Strided(fuel, dl, dr, A, B, ls, rs, Nat.add(ls, lp), Nat.add(rs, rp), values)}
97	
98	# What cell `position` of a build under the layout holds below the count:
99	# the fold, from +0.0 in increasing k, of the `inner` factor pairs that
100	# start at the cell's first reads (`Batched.left_start`, which reads the
101	# left operand's own batch of the cell's batch, a stretched batch axis at
102	# 0, and the cell's row; `Batched.right_start` the same for the right
103	# operand and the cell's column) and step by each operand's stride.
104	type CellAt<-dl: Nat, -dr: Nat, -A: Array<F32>, -B: Array<F32>, -lo: Matrix.Order, -ro: Matrix.Order, -inner: Nat, -rows: Nat,
105	  -columns: Nat, -l1: Nat, -l2: Nat, -r1: Nat, -r2: Nat, -o2: Nat, -position: Nat, -value: Maybe<&2, F32>> is Data:
106	  CellAt{-word: F32, -observed: {value == Some{word} : Maybe<&2, F32>}, -values: List<&2, DotFold.Factors>,
107	    -logical: Strided(inner, dl, dr, A, B, Matrix.stride(lo, rows), Matrix.stride(ro, columns),
108	      Batched.left_start(layout_of(lo, ro, inner, rows, columns, l1, l2, r1, r2, o2), position),
109	      Batched.right_start(layout_of(lo, ro, inner, rows, columns, l1, l2, r1, r2, o2), position), values),
110	    -fold: {word == DotFold.dot_fold(values, 0.0) : F32}}
111	
112	def CellSpec(+dl: Nat, +dr: Nat, -A: Array<F32>, -B: Array<F32>, +lo: Matrix.Order, +ro: Matrix.Order, +inner: Nat, +rows: Nat,
113	  +columns: Nat, +l1: Nat, +l2: Nat, +r1: Nat, +r2: Nat, +o2: Nat, +position: Nat, inside: Bool, -value: Maybe<&2, F32>) -> Type:
114	  match inside:
115	    case False{}:
116	      {value == Some{0.0} : Maybe<&2, F32>}
117	    case True{}:
118	      CellAt<dl, dr, A, B, lo, ro, inner, rows, columns, l1, l2, r1, r2, o2, position, value>
119	
120	# Every slot of `out`, cell `base + offset`, holds what CellSpec names: the
121	# cell's fold below `count`, +0.0 at or above it.
122	def Cells(+depth: Nat, +dl: Nat, +dr: Nat, -A: Array<F32>, -B: Array<F32>, +lo: Matrix.Order, +ro: Matrix.Order, +inner: Nat,
123	  +rows: Nat, +columns: Nat, +l1: Nat, +l2: Nat, +r1: Nat, +r2: Nat, +o2: Nat, +count: Nat, +base: Nat, -out: Array<F32>) -> Type:
124	  @+path: Model.Path(depth) ->
125	    CellSpec(dl, dr, A, B, lo, ro, inner, rows, columns, l1, l2, r1, r2, o2, Nat.add(Model.offset(depth, path), base),
126	      Nat.is_lt(Nat.add(Model.offset(depth, path), base), count), Model.observe(F32, depth, out, path))
127	
128	# For two balanced input arrays A and B and every split, every cell of the
129	# parallel build holds the fold its layout names, so no word depends on the
130	# split or the thread count, and both arrays are only read. The forward
131	# product and both gradients of the tape's BatchedMatmul are such builds
132	# (`Batched.forward_layout`, `left_layout`, `right_layout`).
133	law batched_build_cells:
134	  for +depth: Nat
135	  for +splits: Nat
136	  for +dl: Nat
137	  for +dr: Nat
138	  for -A: Array<F32>
139	  for -B: Array<F32>
140	  for +lo: Matrix.Order
141	  for +ro: Matrix.Order
142	  for +inner: Nat
143	  for +rows: Nat
144	  for +columns: Nat
145	  for +l1: Nat
146	  for +l2: Nat
147	  for +r1: Nat
148	  for +r2: Nat
149	  for +o2: Nat
150	  for +count: Nat
151	  for +base: Nat
152	  for la: Array<F32>
153	  for ra: Array<F32>
154	  for el: {la == A : Array<F32>}
155	  for er: {ra == B : Array<F32>}
156	  for left_valid: Model.Balanced(F32, dl, A)
157	  for right_valid: Model.Balanced(F32, dr, B)
158	  Holds<Cells(depth, dl, dr, A, B, lo, ro, inner, rows, columns, l1, l2, r1, r2, o2, count, base,
159	    Batched.build(depth, splits, layout_of(lo, ro, inner, rows, columns, l1, l2, r1, r2, o2), count, base, la, ra))>
160	
161	# The build is a balanced tree of its depth, whatever the split: the
162	# certificate the output storage carries.
163	law batched_build_balanced:
164	  for +depth: Nat
165	  for +splits: Nat
166	  for +given: Batched.Layout
167	  for +count: Nat
168	  for +base: Nat
169	  for la: Array<F32>
170	  for ra: Array<F32>
171	  Model.Balanced(F32, depth, Batched.build(depth, splits, given, count, base, la, ra))
172	
173	# Which words each cell reads, written from torch.matmul's definition and
174	# the shape module's broadcast, with no call into the batched code.
175	#
176	# Cell (b1, b2, row, column) of an output whose batch has a second extent
177	# of o2 and whose matrices are rows x columns sits at its row-major
178	# position, the cells of one batch after those of the batch before it.
179	def cell_position(+o2: Nat, +rows: Nat, +columns: Nat, +b1: Nat, +b2: Nat, +row: Nat, +column: Nat) -> Nat:
180	  Nat.add(Nat.mul(Nat.add(Nat.mul(b1, o2), b2), Nat.mul(rows, columns)), Nat.add(Nat.mul(row, columns), column))
181	
182	# Term k of row (or column) `outer` of a matrix of `outers` rows (or
183	# columns) and `inner` terms: `[outers, inner]` read by rows (Kept), or
184	# `[inner, outers]` read by columns (Swapped).
185	def matrix_offset(+order: Matrix.Order, +inner: Nat, +outers: Nat, +outer: Nat, +k: Nat) -> Nat:
186	  match order:
187	    case Matrix.Kept{}:
188	      Nat.add(Nat.mul(outer, inner), k)
189	    case Matrix.Swapped{}:
190	      Nat.add(Nat.mul(k, outers), outer)
191	
192	# The row-major offset of that term in an operand of batch extents
193	# [e1, e2] read at batch (b1, b2) of the output: an axis of extent 1 is
194	# broadcast and read at index 0 (`Shape.broadcast_index`, the rule of the
195	# elementwise operations), any other at the output's own index.
196	def operand_offset(+order: Matrix.Order, +inner: Nat, +outers: Nat, +e1: Nat, +e2: Nat, +b1: Nat, +b2: Nat, +outer: Nat, +k: Nat) -> Nat:
197	  Nat.add(Nat.mul(Nat.add(Nat.mul(Shape.broadcast_index(e1, b1), e2), Shape.broadcast_index(e2, b2)), Nat.mul(outers, inner)),
198	    matrix_offset(order, inner, outers, outer, k))
199	
200	# The factor pairs k, k + 1, ... of cell (b1, b2, row, column): the left
201	# operand's word at its term k of row `row`, the right operand's at its
202	# term k of column `column`, each at its own batch.
203	def Logical(+fuel: Nat, +dl: Nat, +dr: Nat, -A: Array<F32>, -B: Array<F32>, +lo: Matrix.Order, +ro: Matrix.Order, +inner: Nat,
204	  +rows: Nat, +columns: Nat, +l1: Nat, +l2: Nat, +r1: Nat, +r2: Nat, +b1: Nat, +b2: Nat, +row: Nat, +column: Nat, +k: Nat,
205	  values: List<&2, DotFold.Factors>) -> Data:
206	  match fuel values:
207	    case 0n Nil{}:
208	      Unit
209	    case 1n+tail DotFold.Factors{lv, rv} <> rest:
210	      LogicalStep<tail, dl, dr, A, B, lo, ro, inner, rows, columns, l1, l2, r1, r2, b1, b2, row, column, k, lv, rv, rest>
211	    case 0n DotFold.Factors{lv, rv} <> rest:
212	      Empty
213	    case 1n+tail Nil{}:
214	      Empty
215	
216	type LogicalStep<-fuel: Nat, -dl: Nat, -dr: Nat, -A: Array<F32>, -B: Array<F32>, -lo: Matrix.Order, -ro: Matrix.Order, -inner: Nat,
217	  -rows: Nat, -columns: Nat, -l1: Nat, -l2: Nat, -r1: Nat, -r2: Nat, -b1: Nat, -b2: Nat, -row: Nat, -column: Nat, -k: Nat,
218	  -lv: F32, -rv: F32, -values: List<&2, DotFold.Factors>> is Data:
219	  LogicalStep{-left_observed: {Model.observe(F32, dl, A, Refinement.route(dl, U32.from_nat(operand_offset(lo, inner, rows, l1, l2, b1, b2, row, k)))) == Some{lv} : Maybe<&2, F32>},
220	    -right_observed: {Model.observe(F32, dr, B, Refinement.route(dr, U32.from_nat(operand_offset(ro, inner, columns, r1, r2, b1, b2, column, k)))) == Some{rv} : Maybe<&2, F32>},
221	    -remaining: Logical(fuel, dl, dr, A, B, lo, ro, inner, rows, columns, l1, l2, r1, r2, b1, b2, row, column, 1n+k, values)}
222	
223	type LogicalCell<-dl: Nat, -dr: Nat, -A: Array<F32>, -B: Array<F32>, -lo: Matrix.Order, -ro: Matrix.Order, -inner: Nat, -rows: Nat,
224	  -columns: Nat, -l1: Nat, -l2: Nat, -r1: Nat, -r2: Nat, -b1: Nat, -b2: Nat, -row: Nat, -column: Nat, -value: Maybe<&2, F32>> is Data:
225	  LogicalCell{-word: F32, -observed: {value == Some{word} : Maybe<&2, F32>}, -values: List<&2, DotFold.Factors>,
226	    -logical: Logical(inner, dl, dr, A, B, lo, ro, inner, rows, columns, l1, l2, r1, r2, b1, b2, row, column, 0n, values),
227	    -fold: {word == DotFold.dot_fold(values, 0.0) : F32}}
228	
229	# For two balanced input arrays and every split, the slot of the build that
230	# holds cell (b1, b2, row, column), with row < rows, column < columns,
231	# b2 < o2 and the cell below the count, holds the fold from +0.0, in
232	# increasing k, of the `inner` factor pairs that `Logical` names: each
233	# operand read at its own batch, an axis of extent 1 at index 0. This is the
234	# meaning of `batched_build_cells` in the output's coordinates; a build
235	# whose reads of a broadcast batch are wrong breaks it. The forward product
236	# reads the left operand by rows and the right one by columns
237	# (`Batched.forward_layout`), the gradients as `left_layout` and
238	# `right_layout` say.
239	law batched_build_logical_cells:
240	  for +depth: Nat
241	  for +splits: Nat
242	  for +dl: Nat
243	  for +dr: Nat
244	  for -A: Array<F32>
245	  for -B: Array<F32>
246	  for +lo: Matrix.Order
247	  for +ro: Matrix.Order
248	  for +inner: Nat
249	  for +rows: Nat
250	  for +columns: Nat
251	  for +l1: Nat
252	  for +l2: Nat
253	  for +r1: Nat
254	  for +r2: Nat
255	  for +o2: Nat
256	  for +count: Nat
257	  for +base: Nat
258	  for la: Array<F32>
259	  for ra: Array<F32>
260	  for el: {la == A : Array<F32>}
261	  for er: {ra == B : Array<F32>}
262	  for left_valid: Model.Balanced(F32, dl, A)
263	  for right_valid: Model.Balanced(F32, dr, B)
264	  for +path: Model.Path(depth)
265	  for +b1: Nat
266	  for +b2: Nat
267	  for +row: Nat
268	  for +column: Nat
269	  for row_below: {Nat.is_lt(row, rows) == True{} : Bool}
270	  for column_below: {Nat.is_lt(column, columns) == True{} : Bool}
271	  for batch_below: {Nat.is_lt(b2, o2) == True{} : Bool}
272	  for at: {Nat.add(Model.offset(depth, path), base) == cell_position(o2, rows, columns, b1, b2, row, column) : Nat}
273	  for inside: {Nat.is_lt(cell_position(o2, rows, columns, b1, b2, row, column), count) == True{} : Bool}
274	  Holds<LogicalCell<dl, dr, A, B, lo, ro, inner, rows, columns, l1, l2, r1, r2, b1, b2, row, column,
275	    Model.observe(F32, depth, Batched.build(depth, splits, layout_of(lo, ro, inner, rows, columns, l1, l2, r1, r2, o2), count, base, la, ra), path)>>
276	
277	# The plan of torch.matmul, written from PyTorch's definition over the
278	# operands' extents as naturals.
279	#
280	# One batch axis of two extents that broadcast: an extent of 1 takes the
281	# other one's.
282	def extent(+left: Nat, +right: Nat) -> Nat:
283	  Bool.pick(Nat, Nat.is_eq(left, 1n), right, left)
284	
285	def vector(dims: List<&2, Nat>) -> Bool:
286	  match dims:
287	    case x <> Nil{}:
288	      True{}
289	    case rest:
290	      False{}
291	
292	# A 1-D left operand [k] is the row [1, k], a 1-D right one the column
293	# [k, 1].
294	def promote_left(dims: List<&2, Nat>) -> List<&2, Nat>:
295	  match dims:
296	    case +x <> Nil{}:
297	      [1n, x]
298	    case rest:
299	      rest
300	
301	def promote_right(dims: List<&2, Nat>) -> List<&2, Nat>:
302	  match dims:
303	    case +x <> Nil{}:
304	      [x, 1n]
305	    case rest:
306	      rest
307	
308	# Every extent but the last two: the batch.
309	def batch_of(dims: List<&2, Nat>) -> List<&2, Nat>:
310	  match dims:
311	    case +x <> +y <> +z <> tail:
312	      x <> batch_of(y <> z <> tail)
313	    case rest:
314	      Nil{}
315	
316	# The extent before the last: the rows of a matrix.
317	def rows_of(dims: List<&2, Nat>) -> Nat:
318	  match dims:
319	    case +x <> +y <> Nil{}:
320	      x
321	    case +x <> +y <> +z <> tail:
322	      rows_of(y <> z <> tail)
323	    case rest:
324	      0n
325	
326	# The last extent: the columns of a matrix.
327	def columns_of(dims: List<&2, Nat>) -> Nat:
328	  match dims:
329	    case +x <> Nil{}:
330	      x
331	    case +x <> +y <> tail:
332	      columns_of(y <> tail)
333	    case Nil{}:
334	      0n
335	
336	def zip_extents(left: List<&2, Nat>, right: List<&2, Nat>) -> List<&2, Nat>:
337	  match left right:
338	    case a <> lt b <> rt:
339	      extent(a, b) <> zip_extents(lt, rt)
340	    case lt rt:
341	      Nil{}
342	
343	# Two batches broadcast right-aligned: the shorter one is padded with ones
344	# on the left (`Shape.pad`), then each axis broadcasts.
345	def broadcast_batch(+left: List<&2, Nat>, +right: List<&2, Nat>) -> List<&2, Nat>:
346	  zip_extents(Shape.pad(Nat.sub(List.length(&2, Nat, right), List.length(&2, Nat, left)), left),
347	    Shape.pad(Nat.sub(List.length(&2, Nat, left), List.length(&2, Nat, right)), right))
348	
349	def kept(dropped: Bool, +extent: Nat) -> List<&2, Nat>:
350	  match dropped:
351	    case True{}:
352	      Nil{}
353	    case False{}:
354	      [extent]
355	
356	# The output: the broadcast batch, then the rows of the promoted left
357	# operand unless it was a vector, then the columns of the promoted right
358	# operand unless it was a vector.
359	def torch_output(+left: List<&2, Nat>, +right: List<&2, Nat>) -> List<&2, Nat>:
360	  List.append(&2, Nat, broadcast_batch(batch_of(promote_left(left)), batch_of(promote_right(right))),
361	    List.append(&2, Nat, kept(vector(left), rows_of(promote_left(left))), kept(vector(right), columns_of(promote_right(right)))))
362	
363	# Extents padded with ones on the left to `rank`.
364	def padded(+rank: Nat, +dims: List<&2, Nat>) -> List<&2, Nat>:
365	  Shape.pad(Nat.sub(rank, List.length(&2, Nat, dims)), dims)
366	
367	# What a plan holds, read back as lists: the left operand as
368	# [l1, l2, rows, inner], the right one as [r1, r2, inner, columns], the
369	# output's batch as [o1, o2], and which operands are vectors.
370	def left_view(+plan: Batched.Plan) -> List<&2, Nat>:
371	  Batched.Plan{output, batch, +rows, +inner, columns, +l1, +l2, r1, r2, o1, o2, lv, rv} = plan
372	  [l1, l2, rows, inner]
373	
374	def right_view(+plan: Batched.Plan) -> List<&2, Nat>:
375	  Batched.Plan{output, batch, rows, +inner, +columns, l1, l2, +r1, +r2, o1, o2, lv, rv} = plan
376	  [r1, r2, inner, columns]
377	
378	def batch_view(+plan: Batched.Plan) -> List<&2, Nat>:
379	  Batched.Plan{output, batch, rows, inner, columns, l1, l2, r1, r2, +o1, +o2, lv, rv} = plan
380	  [o1, o2]
381	
382	def vector_view(+plan: Batched.Plan) -> Bool & Bool:
383	  Batched.Plan{output, batch, rows, inner, columns, l1, l2, r1, r2, o1, o2, +lv, +rv} = plan
384	  (lv, rv)
385	
386	type Geometry<-output: List<&2, Nat>, -left: List<&2, Nat>, -right: List<&2, Nat>, -batch: List<&2, Nat>, -vectors: Bool & Bool,
387	  -plan: Batched.Plan> is Data:
388	  Geometry{-output_dims: {Shape.naturals(Shape.dimensions(Batched.output_of(plan))) == output : List<&2, Nat>},
389	    -left_dims: {left_view(plan) == left : List<&2, Nat>},
390	    -right_dims: {right_view(plan) == right : List<&2, Nat>},
391	    -batch_dims: {batch_view(plan) == batch : List<&2, Nat>},
392	    -vector_dims: {vector_view(plan) == vectors : Bool & Bool}}
393	
394	def GeometryOf(-output: List<&2, Nat>, -left: List<&2, Nat>, -right: List<&2, Nat>, -batch: List<&2, Nat>, -vectors: Bool & Bool,
395	  result: Result<&2, &2, Matrix.Error, Batched.Plan>) -> Data:
396	  match result:
397	    case Fail{error}:
398	      Unit
399	    case Done{plan}:
400	      Geometry<output, left, right, batch, vectors, plan>
401	
402	def torch_geometry(+left: List<&2, Nat>, +right: List<&2, Nat>, result: Result<&2, &2, Matrix.Error, Batched.Plan>) -> Data:
403	  GeometryOf(torch_output(left, right), padded(4n, promote_left(left)), padded(4n, promote_right(right)),
404	    padded(2n, broadcast_batch(batch_of(promote_left(left)), batch_of(promote_right(right)))), (vector(left), vector(right)), result)
405	
406	# Whatever the operands' dimensions, a plan that succeeds has PyTorch's
407	# output dimensions: a 1-D left operand promoted to a row and a 1-D right
408	# one to a column, both dropped again, and the batches broadcast
409	# right-aligned, an extent of 1 taking the other one's. It reads the left
410	# operand as its promoted extents padded with ones to rank 4, the right one
411	# likewise (their inner extents are equal), and the output's batch padded
412	# to two; these are the extents the build's layouts and
413	# `batched_build_logical_cells` read with.
414	law batched_plan_geometry:
415	  for +left: List<&2, U32>
416	  for +right: List<&2, U32>
417	  torch_geometry(Shape.naturals(left), Shape.naturals(right), Batched.plan_dims(left, right))