Support function arange in Array API

.gitignore

@@ -24,3 +24,4 @@ _doc/_static/viz.js
 _unittests/ut__main/*.png
 _unittests/ut__main/_cache/*
 _unittests/ut__main/*.html
+_unittests/.hypothesis/*

_doc/api/index.rst

@@ -8,10 +8,10 @@ API
 
     array_api
     npx_functions
-    npx_var
     npx_jit
-    npx_annot
     npx_numpy
+    npx_types
+    npx_var
     onnx_tools
     ort
     plotting

_doc/api/npx_annot.rst → _doc/api/npx_types.rst

@@ -1,38 +1,40 @@
-=============
 npx.npx_types
 =============
 
 DType
-=====
++++++
 
 .. autoclass:: onnx_array_api.npx.npx_types.DType
     :members:
 
-Annotations
-===========
-
 ElemType
 ++++++++
 
 .. autoclass:: onnx_array_api.npx.npx_types.ElemType
     :members:
 
-ParType
-+++++++
-
-.. autoclass:: onnx_array_api.npx.npx_types.ParType
-    :members:
-
 OptParType
 ++++++++++
 
 .. autoclass:: onnx_array_api.npx.npx_types.OptParType
     :members:
 
-TensorType
-++++++++++
+OptTensorType
++++++++++++++
 
-.. autoclass:: onnx_array_api.npx.npx_types.TensorType
+.. autoclass:: onnx_array_api.npx.npx_types.OptTensorType
+    :members:
+
+ParType
++++++++
+
+.. autoclass:: onnx_array_api.npx.npx_types.ParType
+    :members:
+
+Scalar
+++++++
+
+.. autoclass:: onnx_array_api.npx.npx_types.Scalar
     :members:
 
 SequenceType
@@ -41,6 +43,18 @@ SequenceType
 .. autoclass:: onnx_array_api.npx.npx_types.SequenceType
     :members:
 
+ShapeType
++++++++++
+
+.. autoclass:: onnx_array_api.npx.npx_types.ShapeType
+    :members:
+
+TensorType
+++++++++++
+
+.. autoclass:: onnx_array_api.npx.npx_types.TensorType
+    :members:
+
 TupleType
 +++++++++
 

_doc/api/npx_var.rst

@@ -15,3 +15,16 @@ Cst, Input
 
 .. autoclass:: onnx_array_api.npx.npx_var.Input
     :members:
+
+ManyIdentity
+++++++++++++
+
+.. autoclass:: onnx_array_api.npx.npx_var.ManyIdentity
+    :members:
+
+Par
++++
+
+.. autoclass:: onnx_array_api.npx.npx_var.Par
+    :members:
+

_unittests/onnx-numpy-skips.txt

@@ -1,6 +1,7 @@
 # API failures
 # see https://github.com/data-apis/array-api-tests/blob/master/numpy-skips.txt
-array_api_tests/test_creation_functions.py::test_arange
+array_api_tests/test_creation_functions.py::test_asarray_scalars
+# array_api_tests/test_creation_functions.py::test_arange
 array_api_tests/test_creation_functions.py::test_asarray_arrays
 array_api_tests/test_creation_functions.py::test_empty
 array_api_tests/test_creation_functions.py::test_empty_like

_unittests/test_array_api.sh

@@ -1,4 +1,4 @@
 export ARRAY_API_TESTS_MODULE=onnx_array_api.array_api.onnx_numpy
-pytest ../array-api-tests/array_api_tests/test_creation_functions.py::test_asarray_scalars || exit 1
+pytest ../array-api-tests/array_api_tests/test_creation_functions.py::test_arange || exit 1
 # pytest ../array-api-tests/array_api_tests/test_creation_functions.py --help
-pytest ../array-api-tests/array_api_tests/test_creation_functions.py --hypothesis-explain --skips-file=_unittests/onnx-numpy-skips.txt || exit 1
+pytest ../array-api-tests/array_api_tests/test_creation_functions.py --hypothesis-explain --skips-file=_unittests/onnx-numpy-skips.txt || exit 1

_unittests/ut_array_api/test_onnx_numpy.py

@@ -1,3 +1,4 @@
+import sys
 import unittest
 import numpy as np
 from onnx_array_api.ext_test_case import ExtTestCase
@@ -19,6 +20,22 @@ def test_zeros(self):
         a = xp.absolute(mat)
         self.assertEqualArray(np.absolute(mat.numpy()), a.numpy())
 
+    def test_arange_default(self):
+        a = EagerTensor(np.array([0], dtype=np.int64))
+        b = EagerTensor(np.array([2], dtype=np.int64))
+        mat = xp.arange(a, b)
+        matnp = mat.numpy()
+        self.assertEqual(matnp.shape, (2,))
+        self.assertEqualArray(matnp, np.arange(0, 2).astype(np.int64))
+
+    def test_arange_step(self):
+        a = EagerTensor(np.array([4], dtype=np.int64))
+        s = EagerTensor(np.array([2], dtype=np.int64))
+        mat = xp.arange(a, step=s)
+        matnp = mat.numpy()
+        self.assertEqual(matnp.shape, (2,))
+        self.assertEqualArray(matnp, np.arange(4, step=2).astype(np.int64))
+
     def test_zeros_none(self):
         c = EagerTensor(np.array([4, 5], dtype=np.int64))
         mat = xp.zeros(c)
@@ -52,7 +69,27 @@ def test_full_bool(self):
         self.assertNotEmpty(matnp[0, 0])
         self.assertEqualArray(matnp, np.full((4, 5), False))
 
+    def test_arange_int00a(self):
+        a = EagerTensor(np.array([0], dtype=np.int64))
+        b = EagerTensor(np.array([0], dtype=np.int64))
+        mat = xp.arange(a, b)
+        matnp = mat.numpy()
+        self.assertEqual(matnp.shape, (0,))
+        expected = np.arange(0, 0)
+        if sys.platform == "win32":
+            expected = expected.astype(np.int64)
+        self.assertEqualArray(matnp, expected)
+
+    def test_arange_int00(self):
+        mat = xp.arange(0, 0)
+        matnp = mat.numpy()
+        self.assertEqual(matnp.shape, (0,))
+        expected = np.arange(0, 0)
+        if sys.platform == "win32":
+            expected = expected.astype(np.int64)
+        self.assertEqualArray(matnp, expected)
+
 
 if __name__ == "__main__":
-    TestOnnxNumpy().test_zeros_none()
+    TestOnnxNumpy().test_arange_int00()
     unittest.main(verbosity=2)

_unittests/ut_npx/test_npx.py

@@ -103,6 +103,7 @@
     Int64,
     OptParType,
     TensorType,
+    OptTensorType,
 )
 from onnx_array_api.npx.npx_var import Input, Var
 
@@ -125,35 +126,62 @@ def test_shape_inference(self):
         self.assertEqual(output.type.tensor_type.elem_type, TensorProto.FLOAT)
 
     def test_tensor(self):
-        dt = TensorType["float32"]
+        dt = TensorType["float32", "F32"]
         self.assertEqual(len(dt.dtypes), 1)
         self.assertEqual(dt.dtypes[0].dtype, ElemType.float32)
         self.assertEmpty(dt.shape)
-        self.assertEqual(dt.type_name(), "TensorType['float32']")
+        self.assertEqual(dt.type_name(), "TensorType['float32', 'F32']")
 
-        dt = TensorType["float32"]
+        dt = TensorType["float32", "F32"]
         self.assertEqual(len(dt.dtypes), 1)
         self.assertEqual(dt.dtypes[0].dtype, ElemType.float32)
-        self.assertEqual(dt.type_name(), "TensorType['float32']")
+        self.assertEqual(dt.type_name(), "TensorType['float32', 'F32']")
 
-        dt = TensorType[np.float32]
+        dt = TensorType[np.float32, "F32"]
         self.assertEqual(len(dt.dtypes), 1)
         self.assertEqual(dt.dtypes[0].dtype, ElemType.float32)
-        self.assertEqual(dt.type_name(), "TensorType['float32']")
+        self.assertEqual(dt.type_name(), "TensorType['float32', 'F32']")
         self.assertEmpty(dt.shape)
 
-        dt = TensorType[np.str_]
+        dt = TensorType[np.str_, "TEXT"]
         self.assertEqual(len(dt.dtypes), 1)
         self.assertEqual(dt.dtypes[0].dtype, ElemType.str_)
-        self.assertEqual(dt.type_name(), "TensorType[strings]")
+        self.assertEqual(dt.type_name(), "TensorType[strings, 'TEXT']")
+        self.assertEmpty(dt.shape)
+
+        self.assertRaise(lambda: TensorType[None], TypeError)
+        self.assertRaise(lambda: TensorType[{np.float32, np.str_}], TypeError)
+
+    def test_opt_tensor(self):
+        dt = OptTensorType["float32", "F32"]
+        self.assertEqual(len(dt.dtypes), 1)
+        self.assertEqual(dt.dtypes[0].dtype, ElemType.float32)
+        self.assertEmpty(dt.shape)
+        self.assertEqual(dt.type_name(), "OptTensorType['float32', 'F32']")
+
+        dt = OptTensorType["float32", "F32"]
+        self.assertEqual(len(dt.dtypes), 1)
+        self.assertEqual(dt.dtypes[0].dtype, ElemType.float32)
+        self.assertEqual(dt.type_name(), "OptTensorType['float32', 'F32']")
+
+        dt = OptTensorType[np.float32, "F32"]
+        self.assertEqual(len(dt.dtypes), 1)
+        self.assertEqual(dt.dtypes[0].dtype, ElemType.float32)
+        self.assertEqual(dt.type_name(), "OptTensorType['float32', 'F32']")
+        self.assertEmpty(dt.shape)
+
+        dt = OptTensorType[np.str_, "TEXT"]
+        self.assertEqual(len(dt.dtypes), 1)
+        self.assertEqual(dt.dtypes[0].dtype, ElemType.str_)
+        self.assertEqual(dt.type_name(), "OptTensorType[strings, 'TEXT']")
         self.assertEmpty(dt.shape)
 
         self.assertRaise(lambda: TensorType[None], TypeError)
         self.assertRaise(lambda: TensorType[{np.float32, np.str_}], TypeError)
 
     def test_superset(self):
-        t1 = TensorType[ElemType.numerics]
-        t2 = TensorType[ElemType.float64]
+        t1 = TensorType[ElemType.numerics, "T"]
+        t2 = TensorType[ElemType.float64, "F64"]
         self.assertTrue(t1.issuperset(t2))
         t1 = Float32[None]
         t2 = Float32[None]
@@ -167,14 +195,14 @@ def test_superset(self):
         t1 = Float32["N"]
         t2 = Float32[5]
         self.assertTrue(t1.issuperset(t2))
-        t1 = TensorType[ElemType.int64]
+        t1 = TensorType[ElemType.int64, "I"]
         t2 = Int64[1]
         self.assertTrue(t1.issuperset(t2))
 
     def test_sig(self):
         def local1(
-            x: TensorType[ElemType.floats],
-        ) -> TensorType[ElemType.floats]:
+            x: TensorType[ElemType.floats, "T"],
+        ) -> TensorType[ElemType.floats, "T"]:
             return x
 
         def local2(
@@ -2536,13 +2564,17 @@ def test_numpy_all_empty_axis_1(self):
         got = ref.run(None, {"A": data})
         self.assertEqualArray(y, got[0])
 
-    @unittest.skipIf(True, reason="Fails to follow Array API")
-    def test_get_item(self):
+    def test_get_item_b(self):
         a = EagerNumpyTensor(np.array([True], dtype=np.bool_))
         i = a[0]
         self.assertEqualArray(i.numpy(), a.numpy()[0])
 
+    def test_get_item_i8(self):
+        a = EagerNumpyTensor(np.array([5, 6], dtype=np.int8))
+        i = a[0]
+        self.assertEqualArray(i.numpy(), a.numpy()[0])
+
 
 if __name__ == "__main__":
-    # TestNpx().test_get_item()
+    TestNpx().test_filter()
     unittest.main(verbosity=2)

_unittests/win_test_array_api.bat

@@ -0,0 +1,4 @@
+@echo off
+set ARRAY_API_TESTS_MODULE=onnx_array_api.array_api.onnx_numpy
+python -m pytest ../../array-api-tests/array_api_tests/test_creation_functions.py::test_arange || exit 1
+python -m pytest ../../array-api-tests/array_api_tests/test_creation_functions.py --hypothesis-explain --skips-file=_unittests/onnx-numpy-skips.txt || exit 1

onnx_array_api/array_api/onnx_numpy.py

@@ -3,6 +3,7 @@
 """
 from typing import Any, Optional
 import numpy as np
+from onnx import TensorProto
 from ..npx.npx_functions import (
     all,
     abs,
@@ -15,18 +16,28 @@
     reshape,
     take,
 )
+from ..npx.npx_functions import arange as generic_arange
 from ..npx.npx_functions import full as generic_full
 from ..npx.npx_functions import ones as generic_ones
 from ..npx.npx_functions import zeros as generic_zeros
 from ..npx.npx_numpy_tensors import EagerNumpyTensor
-from ..npx.npx_types import DType, ElemType, TensorType, OptParType, ParType, Scalar
+from ..npx.npx_types import (
+    DType,
+    ElemType,
+    TensorType,
+    OptParType,
+    OptTensorType,
+    ParType,
+    Scalar,
+)
 from ._onnx_common import template_asarray
 from . import _finalize_array_api
 
 __all__ = [
     "abs",
     "absolute",
     "all",
+    "arange",
     "asarray",
     "astype",
     "empty",
@@ -57,6 +68,44 @@ def asarray(
     )
 
 
+def arange(
+    start_or_stop: TensorType[ElemType.int64, "I", (1,)],
+    stop_or_step: OptTensorType[ElemType.int64, "I", (1,)] = None,
+    step: OptTensorType[ElemType.int64, "I", (1,)] = None,
+    dtype: OptParType[DType] = None,
+) -> TensorType[ElemType.numerics, "T"]:
+    use_float = any(
+        map(lambda x: isinstance(x, float), [start_or_stop, stop_or_step, step])
+    )
+    if isinstance(start_or_stop, int):
+        start_or_stop = EagerNumpyTensor(
+            np.array([start_or_stop], dtype=np.float64 if use_float else np.int64)
+        )
+    elif isinstance(start_or_stop, float):
+        start_or_stop = EagerNumpyTensor(np.array([start_or_stop], dtype=np.float64))
+        assert use_float
+
+    if isinstance(stop_or_step, int):
+        stop_or_step = EagerNumpyTensor(
+            np.array([stop_or_step], dtype=np.float64 if use_float else np.int64)
+        )
+    elif isinstance(stop_or_step, float):
+        stop_or_step = EagerNumpyTensor(np.array([stop_or_step], dtype=np.float64))
+        assert use_float
+
+    if isinstance(step, int):
+        step = EagerNumpyTensor(
+            np.array([step], dtype=np.float64 if use_float else np.int64)
+        )
+    elif isinstance(step, float):
+        step = EagerNumpyTensor(np.array([step], dtype=np.float64))
+        assert use_float
+
+    if dtype is None and use_float:
+        dtype = DType(TensorProto.DOUBLE)
+    return generic_arange(start_or_stop, stop_or_step, step, dtype=dtype)
+
+
 def ones(
     shape: TensorType[ElemType.int64, "I", (None,)],
     dtype: OptParType[DType] = None,