Support function arange in Array API (#19)

* add arange

* introduce OptTensorType

* add more tests

* better error handling

* add kwargs_to_input

* fix inconcistencies

* improvments

* fix one type issue

* issue with windows

* set

* remove unnecessary code

* improvments

* fix names

* fix missing name

* fix arange

* fix arange

* fix unit test for windows

---------

Co-authored-by: xavier dupré <[email protected]>

Author: xadupre

.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 renamed to _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)