src/tensor/batched_matmul_LAWS.bend

bendtorch-core@0.1.0.0 · 417 lines · 40 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/matmul.bend as Matrix
7import bendtorch-core@0.1.0.0/src/tensor/batched_matmul.bend as Batched
8import bendtorch-core@0.1.0.0/src/tensor/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
14type 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.
20law 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").
37law 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.
48law 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.
60law 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.
74def 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.
80def 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
92type 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.
104type 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
112def 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.
122def 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`).
133law 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.
163law 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.
179def 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).
185def 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.
196def 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.
203def 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
216type 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
223type 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.
239law 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.
282def extent(+left: Nat, +right: Nat) -> Nat:
283  Bool.pick(Nat, Nat.is_eq(left, 1n), right, left)
284
285def 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].
294def 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
301def 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.
309def 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.
317def 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.
327def 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
336def 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.
345def 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
349def 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.
359def 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`.
364def 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.
370def 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
374def 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
378def 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
382def 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
386type 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
394def 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
402def 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.
414law 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))