fix: support nested block tensor literals [[ta; tb]; [tc; td]]

Add Concat-Concat element-wise unification case to row.ml's unify_dim,
allowing the shape solver to equate structurally-matching Concat dimensions
from inner block tensor results when used in outer block tensors.

Update test and docs to demonstrate direct nesting (2x2 block matrix).

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>

Author: lukstafi

docs/syntax_extensions.md

@@ -574,10 +574,8 @@ will trigger block tensor interpretation.
 let%op stacked = [v1; v2; v3]
 (* Shape: if each vi has output [d], result has output [3, d] *)
 
-(* 2x2 block matrix from scalars -- nesting currently requires let bindings *)
-let%op row1 = [a; b] in
-let%op row2 = [c; d] in
-let%op mat = (row1, row2) ++^ "a; b => a^b"
+(* 2x2 block matrix from vectors via direct nesting *)
+let%op mat = [[a; b]; [c; d]]
 
 (* Batch two samples *)
 let%op batched = [|sample1; sample2|]
@@ -592,8 +590,8 @@ let%op input_block = (x, y)
 - Single-element block tensors like `[ta]` act as unsqueeze (add a size-1 leading axis).
 - Tuple syntax `(ta, tb)` only works at the top level of a `%op` expression. Tuples inside function
   arguments (e.g., `f (ta, tb)`) are preserved as regular OCaml tuples.
-- Direct nesting like `[[ta; tb]; [tc; td]]` is currently limited by shape inference; use
-  intermediate let bindings or explicit `++^` for nested block construction.
+- Nesting works: `[[ta; tb]; [tc; td]]` constructs a 2x2 block matrix (two new output axes).
+  Components at each nesting level must have the same shape.
 
 ### Capturing the dimensions of selected axes for further computation or to add shape constraints
 

tensor/row.ml

@@ -1884,6 +1884,11 @@ let%track6_sexp rec unify_dim ~stage origin (eq : dim * dim) env : constraint_ l
         (more_ineqs @ ineqs, env)
       in
       List.fold ~init:(ineqs, env) dim_eqs ~f
+  | Concat dims1, Concat dims2 when List.length dims1 = List.length dims2 ->
+      (* Unify element-wise: Concat [a; b] = Concat [c; d] iff a = c and b = d *)
+      List.fold2_exn dims1 dims2 ~init:([], env) ~f:(fun (ineqs, env) d1 d2 ->
+          let more_ineqs, env = unify_dim ~stage origin (d1, d2) env in
+          (more_ineqs @ ineqs, env))
   | dim1, dim2 ->
       (* Note: at the unify_dim phase, it's strict equality (no broadcasting). *)
       raise @@ Shape_error ("solved dimensions for axis: mismatch", [ Dim_mismatch [ dim1; dim2 ] ])

test/operations/test_block_tensor.expected

@@ -32,19 +32,22 @@ HERE: test/operations/test_block_tensor.ml:59:21
   ; [ 7.00 ; 8.00 ; 9.00 ]
 ]
 
---- Test 5: Scalars in block tensor ---
-HERE: test/operations/test_block_tensor.ml:73:21
-[48]: ++^_scalar_stack shape 0:3  [ 1.00 ; 2.00 ; 3.00 ]
+--- Test 5: Nested block matrix [[a; b]; [c; d]] ---
+HERE: test/operations/test_block_tensor.ml:78:21
+[60]: ++^_block_matrix shape 0:2,1:2,2:2  [
+  [ [ 1.00 ; 2.00 ] ; [ 3.00 ; 4.00 ] ]
+  ; [ [ 5.00 ; 6.00 ] ; [ 7.00 ; 8.00 ] ]
+]
 
 --- Test 6: Single element [x1] ---
-HERE: test/operations/test_block_tensor.ml:80:21
-[52]: ++^_unsqueezed shape 0:1,1:3  [ [ 1.00 ; 2.00 ; 3.00 ] ]
+HERE: test/operations/test_block_tensor.ml:85:21
+[64]: ++^_unsqueezed shape 0:1,1:3  [ [ 1.00 ; 2.00 ; 3.00 ] ]
 
 --- Test 7: Gradient flow (2-way) ---
 grad_result (sin of stacked):
-HERE: test/operations/test_block_tensor.ml:99:21
+HERE: test/operations/test_block_tensor.ml:104:21
 ┌────────────────────────────────────┐
-│[64]: sin_grad_result shape 0:2,1:2 │
+│[76]: sin_grad_result shape 0:2,1:2 │
 │┌──────┬───────────────────┐        │
 ││      │axis 1             │        │
 │├──────┼───────────────────┤        │
@@ -54,17 +57,17 @@ HERE: test/operations/test_block_tensor.ml:99:21
 └────────────────────────────────────┘
 
 Gradient of g1 (should be cos of original):
-HERE: test/operations/test_block_tensor.ml:101:21
+HERE: test/operations/test_block_tensor.ml:106:21
 ┌────────────────────┐
-│[54]: 1,2 shape 0:2 │
+│[66]: 1,2 shape 0:2 │
 │┌┬────────────┐     │
 │││axis 0      │     │
 │├┼────────────┤     │
 │││ 1.00  2.00 │     │
 │└┴────────────┘     │
 └────────────────────┘
 ┌─────────────────────────────┐
-│[54]: 1,2 shape 0:2  grad_1,2│
+│[66]: 1,2 shape 0:2  grad_1,2│
 │┌┬───────────────────┐       │
 │││axis 0             │       │
 │├┼───────────────────┤       │
@@ -73,17 +76,17 @@ HERE: test/operations/test_block_tensor.ml:101:21
 └─────────────────────────────┘
 
 Gradient of g2:
-HERE: test/operations/test_block_tensor.ml:103:21
+HERE: test/operations/test_block_tensor.ml:108:21
 ┌────────────────────┐
-│[56]: 3,4 shape 0:2 │
+│[68]: 3,4 shape 0:2 │
 │┌┬────────────┐     │
 │││axis 0      │     │
 │├┼────────────┤     │
 │││ 3.00  4.00 │     │
 │└┴────────────┘     │
 └────────────────────┘
 ┌─────────────────────────────┐
-│[56]: 3,4 shape 0:2  grad_3,4│
+│[68]: 3,4 shape 0:2  grad_3,4│
 │┌┬────────────────────┐      │
 │││axis 0              │      │
 │├┼────────────────────┤      │
@@ -93,9 +96,9 @@ HERE: test/operations/test_block_tensor.ml:103:21
 
 --- Test 8: Gradient flow (3-way) ---
 grad3_result (sin of 3-way stacked):
-HERE: test/operations/test_block_tensor.ml:126:21
+HERE: test/operations/test_block_tensor.ml:131:21
 ┌─────────────────────────────────────┐
-│[84]: sin_grad3_result shape 0:3,1:2 │
+│[96]: sin_grad3_result shape 0:3,1:2 │
 │┌──────┬──────────────────┐          │
 ││      │axis 1            │          │
 │├──────┼──────────────────┤          │
@@ -106,17 +109,17 @@ HERE: test/operations/test_block_tensor.ml:126:21
 └─────────────────────────────────────┘
 
 Gradient of h1:
-HERE: test/operations/test_block_tensor.ml:128:21
+HERE: test/operations/test_block_tensor.ml:133:21
 ┌────────────────────────┐
-│[70]: 0.5,1.5 shape 0:2 │
+│[82]: 0.5,1.5 shape 0:2 │
 │┌┬───────────────┐      │
 │││axis 0         │      │
 │├┼───────────────┤      │
 │││ 5.00e-1  1.50 │      │
 │└┴───────────────┘      │
 └────────────────────────┘
 ┌─────────────────────────────────────┐
-│[70]: 0.5,1.5 shape 0:2  grad_0.5,1.5│
+│[82]: 0.5,1.5 shape 0:2  grad_0.5,1.5│
 │┌┬──────────────────┐                │
 │││axis 0            │                │
 │├┼──────────────────┤                │
@@ -125,17 +128,17 @@ HERE: test/operations/test_block_tensor.ml:128:21
 └─────────────────────────────────────┘
 
 Gradient of h2:
-HERE: test/operations/test_block_tensor.ml:130:21
+HERE: test/operations/test_block_tensor.ml:135:21
 ┌────────────────────┐
-│[72]: 1,2 shape 0:2 │
+│[84]: 1,2 shape 0:2 │
 │┌┬────────────┐     │
 │││axis 0      │     │
 │├┼────────────┤     │
 │││ 1.00  2.00 │     │
 │└┴────────────┘     │
 └────────────────────┘
 ┌─────────────────────────────┐
-│[72]: 1,2 shape 0:2  grad_1,2│
+│[84]: 1,2 shape 0:2  grad_1,2│
 │┌┬───────────────────┐       │
 │││axis 0             │       │
 │├┼───────────────────┤       │
@@ -144,17 +147,17 @@ HERE: test/operations/test_block_tensor.ml:130:21
 └─────────────────────────────┘
 
 Gradient of h3:
-HERE: test/operations/test_block_tensor.ml:132:21
+HERE: test/operations/test_block_tensor.ml:137:21
 ┌──────────────────────┐
-│[74]: 3,0.1 shape 0:2 │
+│[86]: 3,0.1 shape 0:2 │
 │┌┬───────────────┐    │
 │││axis 0         │    │
 │├┼───────────────┤    │
 │││ 3.00  1.00e-1 │    │
 │└┴───────────────┘    │
 └──────────────────────┘
 ┌─────────────────────────────────┐
-│[74]: 3,0.1 shape 0:2  grad_3,0.1│
+│[86]: 3,0.1 shape 0:2  grad_3,0.1│
 │┌┬───────────────────┐           │
 │││axis 0             │           │
 │├┼───────────────────┤           │

test/operations/test_block_tensor.ml

@@ -58,19 +58,24 @@ let () =
   let ctx = Train.forward_once ctx triple in
   Train.printf ~here:[%here] ~with_code:false ~with_grad:false ~style:`Inline triple;
 
-  (* --- Test 5: Scalars in block tensor --- *)
-  (* NOTE: Nesting like [[s1; s2]; [s3; s4]] or [row1; row2] where row1/row2 are
-     themselves block tensor results is currently limited by shape inference:
-     the Concat dimension types produced by inner blocks can't be reconciled
-     with the row variable (...) in the outer concat spec. Use explicit ++^ for nesting. *)
-  printf "\n--- Test 5: Scalars in block tensor ---\n%!";
-  let s1 = PDSL.ndarray [| 1.0 |] ~batch_dims:[] ~input_dims:[] ~output_dims:[] () in
-  let s2 = PDSL.ndarray [| 2.0 |] ~batch_dims:[] ~input_dims:[] ~output_dims:[] () in
-  let s3 = PDSL.ndarray [| 3.0 |] ~batch_dims:[] ~input_dims:[] ~output_dims:[] () in
-  let%op scalar_stack = [s1; s2; s3] in
-  Train.set_hosted scalar_stack.value;
-  let ctx = Train.forward_once ctx scalar_stack in
-  Train.printf ~here:[%here] ~with_code:false ~with_grad:false ~style:`Inline scalar_stack;
+  (* --- Test 5: Nested block matrix [[a; b]; [c; d]] --- *)
+  printf "\n--- Test 5: Nested block matrix [[a; b]; [c; d]] ---\n%!";
+  let s1 =
+    PDSL.ndarray [| 1.0; 2.0 |] ~batch_dims:[] ~input_dims:[] ~output_dims:[ 2 ] ()
+  in
+  let s2 =
+    PDSL.ndarray [| 3.0; 4.0 |] ~batch_dims:[] ~input_dims:[] ~output_dims:[ 2 ] ()
+  in
+  let s3 =
+    PDSL.ndarray [| 5.0; 6.0 |] ~batch_dims:[] ~input_dims:[] ~output_dims:[ 2 ] ()
+  in
+  let s4 =
+    PDSL.ndarray [| 7.0; 8.0 |] ~batch_dims:[] ~input_dims:[] ~output_dims:[ 2 ] ()
+  in
+  let%op block_matrix = [[s1; s2]; [s3; s4]] in
+  Train.set_hosted block_matrix.value;
+  let ctx = Train.forward_once ctx block_matrix in
+  Train.printf ~here:[%here] ~with_code:false ~with_grad:false ~style:`Inline block_matrix;
 
   (* --- Test 6: Single element [x1] — unsqueeze --- *)
   printf "\n--- Test 6: Single element [x1] ---\n%!";