Fix concat error (#139)

Author: cbourjau

CHANGELOG.rst

@@ -12,10 +12,11 @@ Changelog
 
 **Bug fixes**
 
+- :func:`ndonnx.concat` no longer raises an error if ``axis=None``, the resulting data type is ``int32`` or ``int64``, and one of the provided arrays is zero-sized.
 - :func:`ndonnx.__array_namespace_info__.capabilities()` now reports the number of supported dimensions via the ``"max dimensions"`` entry rather than ``"max rank"``.
-- Add missing onnxruntime workaround for uint32 inputs to ``ndonnx.min`` and ``ndonnx.max``.
-- Fix array instantiation with ``ndonnx.asarray`` and very large Python integers for ``uint64`` data types.
-- Fix passing an Python scalar as the second argument to ``ndonnx.where``.
+- Add missing onnxruntime workaround for uint32 inputs to :func:`ndonnx.min` and :func:`ndonnx.max`.
+- Fix array instantiation with :func:`ndonnx.asarray` and very large Python integers for ``uint64`` data types.
+- Fix passing an Python scalar as the second argument to :func:`ndonnx.where`.
 
 
 **New features**

ndonnx/_typed_array/onnx.py

@@ -724,21 +724,7 @@ def concat(self, others: Sequence[Self], axis: None | int) -> Self:
         if len(arrays) == 1:
             return arrays[0].copy()
 
-        # It seems that there is currently a bug(?) in the type/shape
-        # inference in ONNX which prohibits us from concatenating
-        # empty 1D int32 and int64 arrays (see test case). We therefor
-        # do some hacky special-casing here for those types and those
-        # types only. Other data types are fine.
-        #
-        # TODO: File upstream bug; this may also be what caused the
-        # segfaults in onnxruntime in the past!
-        if self.dtype in (int32, int64) and self.ndim == 1:
-            dummy_axis = op.const([axis + 1], dtype=np.int64)
-            vars = [op.unsqueeze(a._var, dummy_axis) for a in arrays]
-            var = op.concat(vars, axis=axis)
-            var = op.squeeze(var, dummy_axis)
-        else:
-            var = op.concat([a._var for a in arrays], axis=0 if axis is None else axis)
+        var = op.concat([a._var for a in arrays], axis=0 if axis is None else axis)
         return type(self)(var)
 
     def copy(self) -> Self:

ndonnx/_typed_array/ort_compat.py

@@ -27,7 +27,6 @@
 from spox.opset.ai.onnx.v21 import cast as cast
 from spox.opset.ai.onnx.v21 import ceil as ceil
 from spox.opset.ai.onnx.v21 import compress as compress
-from spox.opset.ai.onnx.v21 import concat as concat
 from spox.opset.ai.onnx.v21 import const as const
 from spox.opset.ai.onnx.v21 import exp as exp  # Only for floats in both standards
 from spox.opset.ai.onnx.v21 import expand as expand
@@ -792,3 +791,29 @@ def einsum(
         mapping=mapping,
         fun_name="einsum",
     )(x1, x2)
+
+
+def concat(
+    inputs: Sequence[Var],
+    *,
+    axis: int,
+) -> Var:
+    # It seems that there is currently a bug(?) in the type/shape
+    # inference in ONNX which prohibits us from concatenating
+    # empty 1D int32 and int64 arrays (see test case). We therefor
+    # do some hacky special-casing here for those types and those
+    # types only. Other data types are fine.
+    #
+    # TODO: File upstream bug; this may also be what caused the
+    # segfaults in onnxruntime in the past!
+    tensors = [el.unwrap_tensor() for el in inputs]
+    # All elements must have the same rank at this point
+    first_shape = tensors[0].shape
+    if first_shape is None:
+        raise ValueError("rank of inputs must be statically known")
+    if tensors[0].dtype in (np.int32, np.int64) and len(first_shape) == 1:
+        dummy_axis = op.const([axis + 1], dtype=np.int64)
+        vars = [op.unsqueeze(el, dummy_axis) for el in inputs]
+        var = op.concat(vars, axis=axis)
+        return op.squeeze(var, dummy_axis)
+    return op.concat(inputs, axis=0 if axis is None else axis)

tests/test_manipulation_functions.py

@@ -0,0 +1,17 @@
+# Copyright (c) QuantCo 2023-2025
+# SPDX-License-Identifier: BSD-3-Clause
+
+import numpy as np
+import pytest
+
+import ndonnx as ndx
+
+
+@pytest.mark.skipif(np.__version__ < "2", reason="NumPy 1.x does not provide 'concat'")
+def test_concat_axis_none_zero_sized():
+    def do(npx):
+        a1 = npx.zeros(shape=(0, 0), dtype=npx.int8)
+        a2 = npx.asarray(0, dtype=npx.int32)
+        return npx.concat([a1, a2], axis=None)
+
+    np.testing.assert_array_equal(do(np), do(ndx))