Zero size matmul mitigation (#102)

Author: adityagoel4512

ndonnx/_opset_extensions.py

@@ -437,7 +437,25 @@ def clip(input: _CoreArray, min: _CoreArray, max: _CoreArray) -> _CoreArray:
 
 @eager_propagate
 def matmul(a: _CoreArray, b: _CoreArray) -> _CoreArray:
-    return _CoreArray(op.matmul(a.var, b.var))
+    # TODO(adityagoel4512): this requires an upstream patch in onnxruntime
+    # onnxruntime goes into UB with zero size inputs
+    (out,) = op.if_(
+        op.equal(op.size(a.var), op.const(0, dtype=np.int64)),
+        then_branch=lambda: [
+            op.const(
+                np.zeros(
+                    (),
+                    dtype=np.result_type(
+                        a.var.unwrap_tensor().dtype, b.var.unwrap_tensor().dtype
+                    ),
+                )
+            ),
+        ],
+        else_branch=lambda: [
+            op.matmul(a.var, b.var),
+        ],
+    )
+    return _CoreArray(out)
 
 
 @eager_propagate

tests/test_core.py

@@ -1067,3 +1067,26 @@ def test_repeat_raises(a, repeats, axis):
 def test_getitem_bool_raises(x, index):
     with pytest.raises(IndexError):
         x[index]
+
+
+@pytest.mark.parametrize(
+    "x, y",
+    [
+        (
+            ndx.asarray([], dtype=ndx.uint8),
+            ndx.asarray([], dtype=ndx.uint8),
+        ),
+        (
+            ndx.asarray([], dtype=ndx.float32),
+            ndx.asarray([], dtype=ndx.int16),
+        ),
+        (
+            ndx.asarray([1, 2, 3], dtype=ndx.uint8),
+            ndx.asarray([1, 2, 3], dtype=ndx.float32),
+        ),
+    ],
+)
+def test_matmul_zero_dims(x, y):
+    ndx_result = x @ y
+    np_result = x.to_numpy() @ y.to_numpy()
+    assert_array_equal(ndx_result.to_numpy(), np_result)