Enables linspace (#30)

* Enables test_asarray_scalars

* Add support for array api linspace

* lint

* improves consistency for linspace

* fix linspace

* disable asarrays_arrays

* fix strategies

* aapi

---------

Co-authored-by: Xavier Dupre <[email protected]>

Author: sdpython

_unittests/onnx-numpy-skips.txt

@@ -1,8 +1,8 @@
 # API failures
 # see https://github.com/data-apis/array-api-tests/blob/master/numpy-skips.txt
 # uses __setitem__
-array_api_tests/test_creation_functions.py::test_asarray_arrays
+# 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
-array_api_tests/test_creation_functions.py::test_linspace
+# array_api_tests/test_creation_functions.py::test_linspace
 array_api_tests/test_creation_functions.py::test_meshgrid

_unittests/test_array_api.sh

@@ -1,4 +1,4 @@
 export ARRAY_API_TESTS_MODULE=onnx_array_api.array_api.onnx_numpy
-pytest -v -rxXfE ../array-api-tests/array_api_tests/test_creation_functions.py::test_eye || exit 1
+pytest -v -rxXfE ../array-api-tests/array_api_tests/test_creation_functions.py::test_full_like || exit 1
 # pytest ../array-api-tests/array_api_tests/test_creation_functions.py --help
 pytest -v -rxXfE ../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_hypothesis_array_api.py

@@ -5,7 +5,7 @@
 import numpy as np
 from operator import mul
 from hypothesis import given
-from onnx_array_api.ext_test_case import ExtTestCase
+from onnx_array_api.ext_test_case import ExtTestCase, ignore_warnings
 from onnx_array_api.array_api import onnx_numpy as onxp
 from hypothesis import strategies
 from hypothesis.extra import array_api
@@ -207,11 +207,58 @@ def fctonx(n_rows, n_cols, kw):
         fctonx()
         self.assertEqual(len(args_onxp), len(args_np))
 
+    @ignore_warnings(UserWarning)
+    def test_square_shared_types(self):
+        dtypes = self.onxps.scalar_dtypes()
+        shared_dtypes = strategies.shared(dtypes, key="dtype")
+
+        def shapes(**kw):
+            kw.setdefault("min_dims", 0)
+            kw.setdefault("min_side", 0)
+            return self.onxps.array_shapes(**kw).filter(
+                lambda shape: prod(i for i in shape if i) < self.MAX_ARRAY_SIZE
+            )
+
+        @strategies.composite
+        def kwargs(draw, **kw):
+            result = {}
+            for k, strat in kw.items():
+                if draw(strategies.booleans()):
+                    result[k] = draw(strat)
+            return result
+
+        @strategies.composite
+        def full_like_fill_values(draw):
+            kw = draw(
+                strategies.shared(
+                    kwargs(dtype=strategies.none() | self.onxps.scalar_dtypes()),
+                    key="full_like_kw",
+                )
+            )
+            dtype = kw.get("dtype", None) or draw(shared_dtypes)
+            return draw(self.onxps.from_dtype(dtype))
+
+        args = []
+        sh = shapes()
+        xa = self.onxps.arrays(dtype=shared_dtypes, shape=sh)
+        fu = full_like_fill_values()
+        kws = strategies.shared(
+            kwargs(dtype=strategies.none() | self.onxps.scalar_dtypes()),
+            key="full_like_kw",
+        )
+
+        @given(x=xa, fill_value=fu, kw=kws)
+        def fctonp(x, fill_value, kw):
+            args.append((x, fill_value, kw))
+
+        fctonp()
+        self.assertEqual(len(args), 100)
+
 
 if __name__ == "__main__":
-    # cl = TestHypothesisArraysApis()
-    # cl.setUpClass()
-    # cl.test_scalar_strategies()
+    cl = TestHypothesisArraysApis()
+    cl.setUpClass()
+    cl.test_square_shared_types()
     # import logging
 
     # logging.basicConfig(level=logging.DEBUG)

_unittests/ut_array_api/test_onnx_numpy.py

@@ -2,10 +2,14 @@
 import unittest
 import numpy as np
 from onnx import TensorProto
-from onnx_array_api.ext_test_case import ExtTestCase
+from onnx_array_api.ext_test_case import ExtTestCase, ignore_warnings
 from onnx_array_api.array_api import onnx_numpy as xp
 from onnx_array_api.npx.npx_types import DType
 from onnx_array_api.npx.npx_numpy_tensors import EagerNumpyTensor as EagerTensor
+from onnx_array_api.npx.npx_functions import linspace as linspace_inline
+from onnx_array_api.npx.npx_types import Float64, Int64
+from onnx_array_api.npx.npx_var import Input
+from onnx_array_api.reference import ExtendedReferenceEvaluator
 
 
 class TestOnnxNumpy(ExtTestCase):
@@ -22,6 +26,7 @@ def test_zeros(self):
         a = xp.absolute(mat)
         self.assertEqualArray(np.absolute(mat.numpy()), a.numpy())
 
+    @ignore_warnings(DeprecationWarning)
     def test_arange_default(self):
         a = EagerTensor(np.array([0], dtype=np.int64))
         b = EagerTensor(np.array([2], dtype=np.int64))
@@ -30,6 +35,7 @@ def test_arange_default(self):
         self.assertEqual(matnp.shape, (2,))
         self.assertEqualArray(matnp, np.arange(0, 2).astype(np.int64))
 
+    @ignore_warnings(DeprecationWarning)
     def test_arange_step(self):
         a = EagerTensor(np.array([4], dtype=np.int64))
         s = EagerTensor(np.array([2], dtype=np.int64))
@@ -78,6 +84,7 @@ def test_full_bool(self):
         self.assertNotEmpty(matnp[0, 0])
         self.assertEqualArray(matnp, np.full((4, 5), False))
 
+    @ignore_warnings(DeprecationWarning)
     def test_arange_int00a(self):
         a = EagerTensor(np.array([0], dtype=np.int64))
         b = EagerTensor(np.array([0], dtype=np.int64))
@@ -89,6 +96,7 @@ def test_arange_int00a(self):
             expected = expected.astype(np.int64)
         self.assertEqualArray(matnp, expected)
 
+    @ignore_warnings(DeprecationWarning)
     def test_arange_int00(self):
         mat = xp.arange(0, 0)
         matnp = mat.numpy()
@@ -160,10 +168,94 @@ def test_eye_k(self):
         got = xp.eye(nr, k=1)
         self.assertEqualArray(expected, got.numpy())
 
+    def test_linspace_int(self):
+        a = EagerTensor(np.array([0], dtype=np.int64))
+        b = EagerTensor(np.array([6], dtype=np.int64))
+        c = EagerTensor(np.array(3, dtype=np.int64))
+        mat = xp.linspace(a, b, c)
+        matnp = mat.numpy()
+        expected = np.linspace(a.numpy(), b.numpy(), c.numpy()).astype(np.int64)
+        self.assertEqualArray(expected, matnp)
+
+    def test_linspace_int5(self):
+        a = EagerTensor(np.array([0], dtype=np.int64))
+        b = EagerTensor(np.array([5], dtype=np.int64))
+        c = EagerTensor(np.array(3, dtype=np.int64))
+        mat = xp.linspace(a, b, c)
+        matnp = mat.numpy()
+        expected = np.linspace(a.numpy(), b.numpy(), c.numpy()).astype(np.int64)
+        self.assertEqualArray(expected, matnp)
+
+    def test_linspace_float(self):
+        a = EagerTensor(np.array([0.5], dtype=np.float64))
+        b = EagerTensor(np.array([5.5], dtype=np.float64))
+        c = EagerTensor(np.array(2, dtype=np.int64))
+        mat = xp.linspace(a, b, c)
+        matnp = mat.numpy()
+        expected = np.linspace(a.numpy(), b.numpy(), c.numpy())
+        self.assertEqualArray(expected, matnp)
+
+    def test_linspace_float_noendpoint(self):
+        a = EagerTensor(np.array([0.5], dtype=np.float64))
+        b = EagerTensor(np.array([5.5], dtype=np.float64))
+        c = EagerTensor(np.array(2, dtype=np.int64))
+        mat = xp.linspace(a, b, c, endpoint=0)
+        matnp = mat.numpy()
+        expected = np.linspace(a.numpy(), b.numpy(), c.numpy(), endpoint=0)
+        self.assertEqualArray(expected, matnp)
+
+    @ignore_warnings((RuntimeWarning, DeprecationWarning))  # division by zero
+    def test_linspace_zero(self):
+        expected = np.linspace(0.0, 0.0, 0, endpoint=False)
+        mat = xp.linspace(0.0, 0.0, 0, endpoint=False)
+        matnp = mat.numpy()
+        self.assertEqualArray(expected, matnp)
+
+    @ignore_warnings((RuntimeWarning, DeprecationWarning))  # division by zero
+    def test_linspace_zero_one(self):
+        expected = np.linspace(0.0, 0.0, 1, endpoint=True)
+
+        f = linspace_inline(Input("start"), Input("stop"), Input("num"))
+        onx = f.to_onnx(
+            constraints={
+                "start": Float64[None],
+                "stop": Float64[None],
+                "num": Int64[None],
+                (0, False): Float64[None],
+            }
+        )
+        ref = ExtendedReferenceEvaluator(onx)
+        got = ref.run(
+            None,
+            {
+                "start": np.array(0, dtype=np.float64),
+                "stop": np.array(0, dtype=np.float64),
+                "num": np.array(1, dtype=np.int64),
+            },
+        )
+        self.assertEqualArray(expected, got[0])
+
+        mat = xp.linspace(0.0, 0.0, 1, endpoint=True)
+        matnp = mat.numpy()
+
+        self.assertEqualArray(expected, matnp)
+
+    def test_slice_minus_one(self):
+        g = EagerTensor(np.array([0.0]))
+        expected = g.numpy()[:-1]
+        got = g[:-1]
+        self.assertEqualArray(expected, got.numpy())
+
+    def test_linspace_bug1(self):
+        expected = np.linspace(16777217.0, 0.0, 1)
+        mat = xp.linspace(16777217.0, 0.0, 1)
+        matnp = mat.numpy()
+        self.assertEqualArray(expected, matnp)
+
 
 if __name__ == "__main__":
     # import logging
 
     # logging.basicConfig(level=logging.DEBUG)
-    TestOnnxNumpy().test_eye()
+    TestOnnxNumpy().test_linspace_float_noendpoint()
     unittest.main(verbosity=2)

_unittests/ut_npx/test_npx.py

@@ -20,7 +20,8 @@
 from onnx.reference import ReferenceEvaluator
 from onnx.shape_inference import infer_shapes
 
-from onnx_array_api.ext_test_case import ExtTestCase
+from onnx_array_api.ext_test_case import ExtTestCase, ignore_warnings
+from onnx_array_api.reference import ExtendedReferenceEvaluator
 from onnx_array_api.npx import ElemType, eager_onnx, jit_onnx
 from onnx_array_api.npx.npx_core_api import (
     cst,
@@ -60,6 +61,7 @@
 from onnx_array_api.npx.npx_functions import hstack as hstack_inline
 from onnx_array_api.npx.npx_functions import identity as identity_inline
 from onnx_array_api.npx.npx_functions import isnan as isnan_inline
+from onnx_array_api.npx.npx_functions import linspace as linspace_inline
 from onnx_array_api.npx.npx_functions import log as log_inline
 from onnx_array_api.npx.npx_functions import log1p as log1p_inline
 from onnx_array_api.npx.npx_functions import matmul as matmul_inline
@@ -1654,6 +1656,7 @@ def test_squeeze(self):
         got = ref.run(None, {"A": x})
         self.assertEqualArray(z, got[0])
 
+    @ignore_warnings(DeprecationWarning)
     def test_squeeze_noaxis(self):
         f = squeeze_inline(copy_inline(Input("A")))
         self.assertIsInstance(f, Var)
@@ -2574,6 +2577,51 @@ def test_get_item_i8(self):
         i = a[0]
         self.assertEqualArray(i.numpy(), a.numpy()[0])
 
+    @ignore_warnings(RuntimeWarning)
+    def test_linspace_big_inline(self):
+        # linspace(5, 0, 1) --> [5] even with endpoint=True
+        f = linspace_inline(Input("A"), Input("B"), Input("C"))
+        self.assertIsInstance(f, Var)
+        onx = f.to_onnx(
+            constraints={
+                0: Int64[None],
+                1: Int64[None],
+                2: Int64[None],
+                (0, False): Int64[None],
+            }
+        )
+
+        start = np.array(16777217.0, dtype=np.float64)
+        stop = np.array(0.0, dtype=np.float64)
+        num = np.array(1, dtype=np.int64)
+        y = np.linspace(start, stop, num)
+        ref = ExtendedReferenceEvaluator(onx)
+        got = ref.run(None, {"A": start, "B": stop, "C": num})
+        self.assertEqualArray(y, got[0])
+
+    @ignore_warnings(RuntimeWarning)
+    def test_linspace_inline(self):
+        # linspace(0, 5, 1)
+        f = linspace_inline(Input("A"), Input("B"), Input("C"))
+        self.assertIsInstance(f, Var)
+        onx = f.to_onnx(
+            constraints={
+                0: Int64[None],
+                1: Int64[None],
+                2: Int64[None],
+                (0, False): Int64[None],
+            }
+        )
+
+        start = np.array(0, dtype=np.float64)
+        stop = np.array(5, dtype=np.float64)
+        num = np.array(1, dtype=np.int64)
+        y = np.linspace(start, stop, num)
+        ref = ExtendedReferenceEvaluator(onx)
+        got = ref.run(None, {"A": start, "B": stop, "C": num})
+        self.assertEqualArray(y, got[0])
+
 
 if __name__ == "__main__":
+    TestNpx().test_linspace_inline()
     unittest.main(verbosity=2)

onnx_array_api/array_api/__init__.py

@@ -24,6 +24,7 @@
     "isfinite",
     "isinf",
     "isnan",
+    "linspace",
     "ones",
     "ones_like",
     "reshape",
@@ -40,11 +41,18 @@ def _finfo(dtype):
     """
     dt = dtype.np_dtype if isinstance(dtype, DType) else dtype
     res = np.finfo(dt)
-    d = res.__dict__.copy()
+    d = {}
+    for k, v in res.__dict__.items():
+        if k.startswith("__"):
+            continue
+        if isinstance(v, (np.float32, np.float64, np.float16)):
+            d[k] = float(v)
+        else:
+            d[k] = v
     d["dtype"] = DType(np_dtype_to_tensor_dtype(dt))
     nres = type("finfo", (res.__class__,), d)
-    setattr(nres, "smallest_normal", res.smallest_normal)
-    setattr(nres, "tiny", res.tiny)
+    setattr(nres, "smallest_normal", float(res.smallest_normal))
+    setattr(nres, "tiny", float(res.tiny))
     return nres
 
 
@@ -54,11 +62,30 @@ def _iinfo(dtype):
     """
     dt = dtype.np_dtype if isinstance(dtype, DType) else dtype
     res = np.iinfo(dt)
-    d = res.__dict__.copy()
+    d = {}
+    for k, v in res.__dict__.items():
+        if k.startswith("__"):
+            continue
+        if isinstance(
+            v,
+            (
+                np.int16,
+                np.int32,
+                np.int64,
+                np.uint16,
+                np.uint32,
+                np.uint64,
+                np.int8,
+                np.uint8,
+            ),
+        ):
+            d[k] = int(v)
+        else:
+            d[k] = v
     d["dtype"] = DType(np_dtype_to_tensor_dtype(dt))
-    nres = type("finfo", (res.__class__,), d)
-    setattr(nres, "min", res.min)
-    setattr(nres, "max", res.max)
+    nres = type("iinfo", (res.__class__,), d)
+    setattr(nres, "min", int(res.min))
+    setattr(nres, "max", int(res.max))
     return nres