Fix ndonnx.roll and run more ci examples

Author: cbourjau

.github/workflows/array-api.yml

@@ -1,4 +1,4 @@
-name: Array API coverage tests
+name: Array API tests
 
 on:
   # We would like to trigger for CI for any pull request action -
@@ -28,9 +28,15 @@ jobs:
     strategy:
       fail-fast: false
       matrix:
-        os:
-          - ubuntu-latest
-          - macos-latest
+        include:
+          # Different platforms are faster than others. We try to keep
+          # the CI snappy at <= 2min.
+          - os: ubuntu-latest
+            max_examples_ci: 50
+          - os: macos-latest
+            max_examples_ci: 100
+          - os: windows-latest
+            max_examples_ci: 30
     steps:
       - name: Checkout branch
         uses: actions/checkout@11bd71901bbe5b1630ceea73d27597364c9af683 # v4.2.2
@@ -42,11 +48,12 @@ jobs:
       - name: Install repository
         run: pixi run postinstall
       - name: Run Array API tests (Scheduled)
+        # TODO: Move schedule part into its own workflow
         if: ${{ github.event_name == 'schedule' && github.ref == 'refs/heads/main' }}
         run: pixi run arrayapitests --max-examples 1000 --hypothesis-seed=""
       - name: Run Array API tests (PR and main branch)
         if: ${{ github.event_name != 'schedule' }}
-        run: pixi run arrayapitests
+        run: pixi run arrayapitests --max-examples ${{ matrix.max_examples_ci }}
       - name: Upload Array API tests report
         uses: actions/upload-artifact@ea165f8d65b6e75b540449e92b4886f43607fa02 # v4.6.2
         with:

CHANGELOG.rst

@@ -14,9 +14,10 @@ Changelog
 
 - :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 :func:`ndonnx.min` and :func:`ndonnx.max`.
+- 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`.
+- Calling :func:`ndonnx.roll` on zero-sized inputs no longer causes a segfault on Linux.
 
 
 **New features**

ndonnx/_typed_array/typed_array.py

@@ -12,10 +12,9 @@
 from typing_extensions import Self
 
 from ndonnx import DType
+from ndonnx._typed_array import safe_cast
 from ndonnx.types import OnnxShape, PyScalar
 
-from .ort_compat import const, if_
-
 if TYPE_CHECKING:
     from spox import Var
 
@@ -210,7 +209,15 @@ def roll(
         *,
         axis: int | tuple[int, ...] | None = None,
     ) -> Self:
+        # From a user's perspective this is a dtype-independent
+        # default that is very nice to have. However, some
+        # implementation details require us to use int64 and float64
+        # casts on indices explicitly.
+        from . import onnx
+        from .funcs import where
+
         x = self
+
         axis_ = axis
         if isinstance(shift, int):
             shift = (shift,)
@@ -223,27 +230,26 @@ def roll(
             raise ValueError("'shift' and 'axis' must be tuples of equal length")
 
         def _roll_axis(x: Self, shift: int, axis: int, /) -> Self:
-            from .funcs import astyarray
-
             if shift == 0:
                 return x
 
-            indices_a = [slice(None) for i in range(x.ndim)]
-            indices_b = [slice(None) for i in range(x.ndim)]
-
-            dim = x.dynamic_shape[axis]
-            (shift_,) = map(
-                astyarray,
-                if_(
-                    (dim == 0).disassemble(),
-                    then_branch=lambda: (const(0, dtype=np.int64),),
-                    else_branch=lambda: ((astyarray(shift) % dim).disassemble(),),
-                ),
+            indices_a = [slice(None) for _ in range(x.ndim)]
+            indices_b = [slice(None) for _ in range(x.ndim)]
+
+            # May be zero. mod may crash for integer data types which
+            # is why we compute the following for floats, then fill
+            # the NaNs, and cast back to int
+            dim_length = x.dynamic_shape[axis].astype(onnx.float64)
+            length_mod = safe_cast(
+                onnx.TyArrayFloat64, onnx.const(float(shift), onnx.float64) % dim_length
+            )
+            shift_ = where(length_mod.isnan(), onnx.const(0.0), length_mod).astype(
+                onnx.int64
             )
+
             # pre roll: |----a------|---b---|
             # postroll: |---b---|----a------|
             #           |-shift-|
-
             indices_a[axis] = slice(None, -shift_, 1)
             indices_b[axis] = slice(-shift_, None, 1)
             return x[tuple(indices_b)].concat([x[tuple(indices_a)]], axis=axis)

pixi.toml

@@ -51,7 +51,7 @@ test-coverage = "pytest --cov=ndonnx --cov-report=xml --cov-report=term-missing"
 
 [feature.test.tasks.arrayapitests]
 # Seed and max-examples can be overridden by providing them a second time
-cmd = "pytest array-api-tests/array_api_tests/ -v -rfX --json-report --json-report-file=array-api-tests.json --disable-deadline --disable-extension linalg,fft --skips-file=skips.txt -s --hypothesis-seed 0 --max-examples 16"
+cmd = "pytest array-api-tests/array_api_tests/ -v -rfX --json-report --json-report-file=array-api-tests.json --disable-deadline --disable-extension linalg,fft --skips-file=skips.txt -s --hypothesis-seed 0 --max-examples 100 -n auto"
 [feature.test.tasks.arrayapitests.env]
 ARRAY_API_TESTS_MODULE="ndonnx"
 ARRAY_API_TESTS_VERSION="2024.12"