Fix modulus (#127)

Author: adityagoel4512

.github/workflows/array-api.yml

@@ -23,8 +23,13 @@ jobs:
   array-api-tests:
     # Run if the commit message contains 'run array-api tests' or if the job is triggered on schedule
     name: Array API test
-    timeout-minutes: 90
-    runs-on: ubuntu-latest
+    timeout-minutes: 60
+    runs-on: ${{ matrix.os }}
+    strategy:
+      matrix:
+        os:
+          - ubuntu-latest
+          - macos-latest
     steps:
       - name: Checkout branch
         uses: actions/checkout@11bd71901bbe5b1630ceea73d27597364c9af683 # v4.2.2
@@ -44,8 +49,8 @@ jobs:
       - name: Upload Array API tests report
         uses: actions/upload-artifact@ea165f8d65b6e75b540449e92b4886f43607fa02 # v4.6.2
         with:
-          name: array-api-tests
-          path: array-api-tests.json
+          name: array-api-tests-${{ matrix.os }}
+          path: array-api-tests-${{ matrix.os }}.json
       - name: Issue on failure
         uses: actions/github-script@60a0d83039c74a4aee543508d2ffcb1c3799cdea # v7.0.1
         if: ${{ failure() && github.event_name == 'schedule' && github.ref == 'refs/heads/main' }}

ndonnx/_typed_array/onnx.py

@@ -1577,12 +1577,12 @@ def __or__(self, other) -> TyArrayBase:
     def __ror__(self, other) -> TyArrayBase:
         return self._apply_int_only(other, op.bitwise_or, forward=False)
 
-    def __mod__(self, other) -> TyArrayBase:
+    def __mod__(self, other) -> TyArrayInteger:
         return self._apply_int_only(
             other, lambda a, b: op.mod(a, b, fmod=0), forward=True
         )
 
-    def __rmod__(self, other) -> TyArrayBase:
+    def __rmod__(self, other) -> TyArrayInteger:
         return self._apply_int_only(
             other, lambda a, b: op.mod(a, b, fmod=0), forward=False
         )

ndonnx/_typed_array/typed_array.py

@@ -14,6 +14,8 @@
 from ndonnx import DType
 from ndonnx.types import OnnxShape, PyScalar
 
+from .ort_compat import const, if_
+
 if TYPE_CHECKING:
     from spox import Var
 
@@ -230,7 +232,14 @@ def _roll_axis(x: Self, shift: int, axis: int, /) -> Self:
             indices_b = [slice(None) for i in range(x.ndim)]
 
             dim = x.dynamic_shape[axis]
-            shift_ = astyarray(shift) % dim
+            (shift_,) = map(
+                astyarray,
+                if_(
+                    (dim == 0).disassemble(),
+                    then_branch=lambda: (const(0, dtype=np.int64),),
+                    else_branch=lambda: ((astyarray(shift) % dim).disassemble(),),
+                ),
+            )
             # pre roll: |----a------|---b---|
             # postroll: |---b---|----a------|
             #           |-shift-|