Track 2: 2D-aware broadcasting on elementwise array ops

RandomCoder-lab · claude · RandomCoder-lab · commit 162c67e6741b · 2026-05-16T02:48:25.000-05:00
arr_add / arr_sub / arr_mul / arr_div_int now broadcast 2D-against-2D
(same shape, cell-wise) and 2D-against-1D row-vector (replicate the
vector across every row). The 1D / scalar paths are unchanged, so
existing callers see no behaviour change.

Before this commit, feeding a 2D array to arr_add silently produced
zeros: each row Value coerced via to_int() → 0 in the flat path.
Replacing that with a proper broadcast is strictly more correct.

Shape mismatches raise as typed errors, catchable with try/except
just like Python's ValueError.

The dense-layer test demonstrates the natural composition:
    Z = arr_matmul(X, W)
    Y = arr_add(Z, bias)     # bias as 1D vector, broadcast per row

Tests: 9 cases — 2D+2D add/sub/mul, 2D+1D row broadcast (both
orderings), composition into a dense-layer forward pass, and shape
mismatches caught via the typed-exception path.

Co-Authored-By: Claude Opus 4.7 &lt;noreply@anthropic.com&gt;
diff --git a/examples/tests/test_broadcasting.omc b/examples/tests/test_broadcasting.omc
@@ -0,0 +1,144 @@
+# Track 2: 2D-aware broadcasting on arr_add / arr_sub / arr_mul.
+#
+# Previous behaviour:
+#   - same-length 1D arrays: element-wise
+#   - scalar ↔ 1D: scalar broadcast
+#   - everything else: silently wrong (each row coerced via to_int → 0)
+#
+# Now adds:
+#   - 2D ↔ 2D (same shape): element-wise per cell
+#   - 2D ↔ 1D row-vector: vector replicated across every row
+#   - 1D row-vector ↔ 2D: same, reversed
+
+fn assert_eq(actual, expected, msg) {
+    if actual != expected {
+        test_record_failure(msg + ": expected " + to_string(expected) + " got " + to_string(actual));
+    }
+}
+
+fn assert_true(cond, msg) {
+    if !cond { test_record_failure(msg); }
+}
+
+# ---- 2D + 2D same shape ----
+
+fn test_2d_2d_add() {
+    h A = [[1, 2, 3], [4, 5, 6]];
+    h B = [[10, 20, 30], [40, 50, 60]];
+    h C = arr_add(A, B);
+    assert_eq(arr_len(C), 2, "2 rows");
+    h r0 = arr_get(C, 0);
+    h r1 = arr_get(C, 1);
+    assert_eq(arr_len(r0), 3, "3 cols");
+    assert_eq(arr_get(r0, 0), 11, "(0,0) = 11");
+    assert_eq(arr_get(r0, 2), 33, "(0,2) = 33");
+    assert_eq(arr_get(r1, 0), 44, "(1,0) = 44");
+    assert_eq(arr_get(r1, 2), 66, "(1,2) = 66");
+}
+
+fn test_2d_2d_sub() {
+    h A = [[10, 20], [30, 40]];
+    h B = [[1, 2], [3, 4]];
+    h C = arr_sub(A, B);
+    h r0 = arr_get(C, 0);
+    assert_eq(arr_get(r0, 0), 9, "10-1");
+    assert_eq(arr_get(r0, 1), 18, "20-2");
+    h r1 = arr_get(C, 1);
+    assert_eq(arr_get(r1, 0), 27, "30-3");
+    assert_eq(arr_get(r1, 1), 36, "40-4");
+}
+
+fn test_2d_2d_mul() {
+    h A = [[1, 2], [3, 4]];
+    h B = [[5, 6], [7, 8]];
+    h C = arr_mul(A, B);
+    h r0 = arr_get(C, 0);
+    assert_eq(arr_get(r0, 0), 5, "1*5");
+    assert_eq(arr_get(r0, 1), 12, "2*6");
+    h r1 = arr_get(C, 1);
+    assert_eq(arr_get(r1, 0), 21, "3*7");
+    assert_eq(arr_get(r1, 1), 32, "4*8");
+}
+
+# ---- 2D + 1D row broadcast ----
+
+fn test_2d_1d_add_row_broadcast() {
+    h A = [[1, 2, 3], [10, 20, 30], [100, 200, 300]];
+    h bias = [1, 2, 3];
+    h C = arr_add(A, bias);
+    h r0 = arr_get(C, 0);
+    h r1 = arr_get(C, 1);
+    h r2 = arr_get(C, 2);
+    assert_eq(arr_get(r0, 0), 2, "1+1");
+    assert_eq(arr_get(r0, 2), 6, "3+3");
+    assert_eq(arr_get(r1, 0), 11, "10+1");
+    assert_eq(arr_get(r1, 2), 33, "30+3");
+    assert_eq(arr_get(r2, 0), 101, "100+1");
+    assert_eq(arr_get(r2, 2), 303, "300+3");
+}
+
+fn test_1d_2d_add_symmetric() {
+    h A = [[1, 2, 3], [4, 5, 6]];
+    h bias = [10, 20, 30];
+    # 1D on the LEFT this time
+    h C = arr_add(bias, A);
+    h r0 = arr_get(C, 0);
+    h r1 = arr_get(C, 1);
+    assert_eq(arr_get(r0, 0), 11, "10+1");
+    assert_eq(arr_get(r0, 2), 33, "30+3");
+    assert_eq(arr_get(r1, 1), 25, "20+5");
+}
+
+# ---- Composition: matmul + bias broadcast (a "dense layer") ----
+
+fn test_dense_layer_forward() {
+    # Single example through a linear layer with bias.
+    # X (1x3) @ W (3x2) + b (2,)  = (1x2)
+    h X = [[1, 2, 3]];
+    h W = [[1, 0], [0, 1], [1, 1]];
+    h b = [10, 20];
+    h Z = arr_matmul(X, W);     # (1x2) = [[1+3, 2+3]] = [[4, 5]]
+    h Y = arr_add(Z, b);        # broadcast b across the one row
+    h r0 = arr_get(Y, 0);
+    assert_true(arr_get(r0, 0) > 13, "first activation > 13");
+    assert_true(arr_get(r0, 1) > 24, "second activation > 24");
+}
+
+# ---- Scalar broadcasting still works (regression) ----
+
+fn test_scalar_into_2d_still_unsupported() {
+    # Currently scalar broadcast applies only to 1D arrays; 2D + scalar
+    # would need a separate path. Document the current behaviour so a
+    # future commit that adds it has a baseline to reverse.
+    h A = [[1, 2], [3, 4]];
+    # arr_scale stays a separate builtin and works fine.
+    h C = arr_scale([1, 2, 3], 10);
+    assert_eq(arr_get(C, 0), 10, "scalar broadcast on 1D");
+    assert_eq(arr_get(C, 2), 30, "scalar broadcast on 1D");
+}
+
+# ---- Shape mismatch raises ----
+
+fn test_shape_mismatch_2d_2d_caught() {
+    h caught = "";
+    try {
+        h A = [[1, 2, 3], [4, 5, 6]];
+        h B = [[1, 2], [3, 4]];
+        h _ = arr_add(A, B);
+    } catch e {
+        caught = e;
+    }
+    assert_true(str_len(caught) > 0, "shape mismatch raised");
+}
+
+fn test_shape_mismatch_2d_1d_caught() {
+    h caught = "";
+    try {
+        h A = [[1, 2, 3], [4, 5, 6]];
+        h bad = [1, 2];   # length 2 but 3 cols needed
+        h _ = arr_add(A, bad);
+    } catch e {
+        caught = e;
+    }
+    assert_true(str_len(caught) > 0, "row-broadcast mismatch raised");
+}
diff --git a/omnimcode-core/src/interpreter.rs b/omnimcode-core/src/interpreter.rs
@@ -8337,12 +8337,146 @@ fn values_equal(a: &Value, b: &Value) -> bool {
 /// `op` takes (i64, i64) and returns i64; the helper wraps the
 /// result in HInt so per-element substrate resonance gets recomputed
 /// from the arithmetic output.
+/// Detect whether `a` is a 2D array (every element is itself an array).
+/// Empty rows count as malformed and return None — callers fall back to
+/// the 1D path. Returns (rows, cols) of the first row when 2D.
+fn array_2d_shape(v: &Value) -> Option<(usize, usize)> {
+    if let Value::Array(outer) = v {
+        let rows = outer.items.borrow();
+        if rows.is_empty() { return None; }
+        let first_cols = match &rows[0] {
+            Value::Array(r) => r.items.borrow().len(),
+            _ => return None,
+        };
+        for r in rows.iter() {
+            match r {
+                Value::Array(row) if row.items.borrow().len() == first_cols => {}
+                _ => return None,
+            }
+        }
+        Some((rows.len(), first_cols))
+    } else {
+        None
+    }
+}
+
+/// 2D-aware broadcast paths for elementwise ops. Returns Some(result)
+/// when both operands fit one of the broadcasting shapes; None lets the
+/// caller fall through to the flat 1D path.
+///
+///   (NxM, NxM)        — element-wise, returns NxM
+///   (NxM, M-vector)   — row broadcast: vector added to every row
+///   (M-vector, NxM)   — same, reversed
+fn try_2d_broadcast<F: Fn(i64, i64) -> i64>(
+    a: &Value,
+    b: &Value,
+    name: &str,
+    op: &F,
+) -> Result<Option<Value>, String> {
+    let a_shape = array_2d_shape(a);
+    let b_shape = array_2d_shape(b);
+
+    // Case 1: both 2D — must match shapes element-wise.
+    if let (Some((ar, ac)), Some((br, bc))) = (a_shape, b_shape) {
+        if ar != br || ac != bc {
+            return Err(format!(
+                "{}: 2D shape mismatch ({}x{} vs {}x{})", name, ar, ac, br, bc
+            ));
+        }
+        if let (Value::Array(a_rows), Value::Array(b_rows)) = (a, b) {
+            let ar_b = a_rows.items.borrow();
+            let br_b = b_rows.items.borrow();
+            let mut out_rows: Vec<Value> = Vec::with_capacity(ar);
+            for (ra, rb) in ar_b.iter().zip(br_b.iter()) {
+                let (Value::Array(ra), Value::Array(rb)) = (ra, rb) else {
+                    return Ok(None);
+                };
+                let raw_a = ra.items.borrow();
+                let raw_b = rb.items.borrow();
+                let row: Vec<Value> = raw_a.iter().zip(raw_b.iter())
+                    .map(|(x, y)| Value::HInt(HInt::new(op(x.to_int(), y.to_int()))))
+                    .collect();
+                out_rows.push(Value::Array(HArray::from_vec(row)));
+            }
+            return Ok(Some(Value::Array(HArray::from_vec(out_rows))));
+        }
+    }
+
+    // Case 2: 2D + 1D row-vector — broadcast vector across every row.
+    if let (Some((ar, ac)), None) = (a_shape, b_shape) {
+        if let (Value::Array(a_rows), Value::Array(b_vec)) = (a, b) {
+            let vec_b = b_vec.items.borrow();
+            // Reject when b is itself a non-1D shape (e.g., array of dicts);
+            // a true 1D vector has length == ac.
+            if vec_b.len() != ac {
+                // Could be a length mismatch — surface a clear error.
+                // But only when b looks like a 1D numeric vector; otherwise
+                // fall through to None and let the caller handle.
+                if vec_b.iter().any(|v| matches!(v, Value::Array(_))) {
+                    return Ok(None);
+                }
+                return Err(format!(
+                    "{}: row-broadcast length mismatch ({} cols vs {} vec)",
+                    name, ac, vec_b.len()
+                ));
+            }
+            let ar_b = a_rows.items.borrow();
+            let mut out_rows: Vec<Value> = Vec::with_capacity(ar);
+            for ra in ar_b.iter() {
+                let Value::Array(ra) = ra else { return Ok(None); };
+                let raw_a = ra.items.borrow();
+                let row: Vec<Value> = raw_a.iter().zip(vec_b.iter())
+                    .map(|(x, y)| Value::HInt(HInt::new(op(x.to_int(), y.to_int()))))
+                    .collect();
+                out_rows.push(Value::Array(HArray::from_vec(row)));
+            }
+            return Ok(Some(Value::Array(HArray::from_vec(out_rows))));
+        }
+    }
+
+    // Case 3: 1D + 2D — symmetric.
+    if let (None, Some((br, bc))) = (a_shape, b_shape) {
+        if let (Value::Array(a_vec), Value::Array(b_rows)) = (a, b) {
+            let vec_a = a_vec.items.borrow();
+            if vec_a.len() != bc {
+                if vec_a.iter().any(|v| matches!(v, Value::Array(_))) {
+                    return Ok(None);
+                }
+                return Err(format!(
+                    "{}: row-broadcast length mismatch ({} vec vs {} cols)",
+                    name, vec_a.len(), bc
+                ));
+            }
+            let br_b = b_rows.items.borrow();
+            let mut out_rows: Vec<Value> = Vec::with_capacity(br);
+            for rb in br_b.iter() {
+                let Value::Array(rb) = rb else { return Ok(None); };
+                let raw_b = rb.items.borrow();
+                let row: Vec<Value> = vec_a.iter().zip(raw_b.iter())
+                    .map(|(x, y)| Value::HInt(HInt::new(op(x.to_int(), y.to_int()))))
+                    .collect();
+                out_rows.push(Value::Array(HArray::from_vec(row)));
+            }
+            return Ok(Some(Value::Array(HArray::from_vec(out_rows))));
+        }
+    }
+
+    Ok(None)
+}
+
 pub(crate) fn elementwise_op<F: Fn(i64, i64) -> i64>(
     a: &Value,
     b: &Value,
     name: &str,
     op: F,
 ) -> Result<Value, String> {
+    // 2D-aware broadcasting shortcut — runs before the standard flat-array
+    // path so callers don't have to switch to a separate builtin. Two
+    // 2D operands element-wise; (2D, 1D) row-broadcast (the 1D vector
+    // gets added to every row); (1D, 2D) same in reverse.
+    if let Some(out) = try_2d_broadcast(a, b, name, &op)? {
+        return Ok(out);
+    }
     match (a, b) {
         (Value::Array(arr_a), Value::Array(arr_b)) => {
             let ai = arr_a.items.borrow();
