Supports function full for the Array API (#21)

xadupre · web-flow · commit c82f9f36630e · 2023-06-18T23:06:03.000+02:00
* Supports function full for the Array API

* improvments

* fix keys by adding types

* fix unit tests

* ci
diff --git a/_unittests/onnx-numpy-skips.txt b/_unittests/onnx-numpy-skips.txt
@@ -5,7 +5,6 @@ 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_eye
-array_api_tests/test_creation_functions.py::test_full
 array_api_tests/test_creation_functions.py::test_full_like
 array_api_tests/test_creation_functions.py::test_linspace
 array_api_tests/test_creation_functions.py::test_meshgrid
diff --git a/_unittests/test_array_api.sh b/_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_arrays || exit 1
+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 --help
 pytest ../array-api-tests/array_api_tests/test_creation_functions.py --hypothesis-explain --skips-file=_unittests/onnx-numpy-skips.txt || exit 1
diff --git a/_unittests/ut_array_api/test_array_apis.py b/_unittests/ut_array_api/test_array_apis.py
@@ -13,7 +13,7 @@ class TestArraysApis(ExtTestCase):
     def test_zeros_numpy_1(self):
         c = xpn.zeros(1)
         d = c.numpy()
-        self.assertEqualArray(np.array([0], dtype=np.float32), d)
+        self.assertEqualArray(np.array([0], dtype=np.float64), d)
 
     def test_zeros_ort_1(self):
         c = xpo.zeros(1)
diff --git a/_unittests/ut_array_api/test_onnx_numpy.py b/_unittests/ut_array_api/test_onnx_numpy.py
@@ -19,6 +19,40 @@ def test_zeros(self):
         a = xp.absolute(mat)
         self.assertEqualArray(np.absolute(mat.numpy()), a.numpy())
 
+    def test_zeros_none(self):
+        c = EagerTensor(np.array([4, 5], dtype=np.int64))
+        mat = xp.zeros(c)
+        matnp = mat.numpy()
+        self.assertEqual(matnp.shape, (4, 5))
+        self.assertNotEmpty(matnp[0, 0])
+        self.assertEqualArray(matnp, np.zeros((4, 5)))
+
+    def test_ones_none(self):
+        c = EagerTensor(np.array([4, 5], dtype=np.int64))
+        mat = xp.ones(c)
+        matnp = mat.numpy()
+        self.assertEqual(matnp.shape, (4, 5))
+        self.assertNotEmpty(matnp[0, 0])
+        self.assertEqualArray(matnp, np.ones((4, 5)))
+
+    def test_full(self):
+        c = EagerTensor(np.array([4, 5], dtype=np.int64))
+        mat = xp.full(c, fill_value=5, dtype=xp.int64)
+        matnp = mat.numpy()
+        self.assertEqual(matnp.shape, (4, 5))
+        self.assertNotEmpty(matnp[0, 0])
+        a = xp.absolute(mat)
+        self.assertEqualArray(np.absolute(mat.numpy()), a.numpy())
+
+    def test_full_bool(self):
+        c = EagerTensor(np.array([4, 5], dtype=np.int64))
+        mat = xp.full(c, fill_value=False)
+        matnp = mat.numpy()
+        self.assertEqual(matnp.shape, (4, 5))
+        self.assertNotEmpty(matnp[0, 0])
+        self.assertEqualArray(matnp, np.full((4, 5), False))
+
 
 if __name__ == "__main__":
+    TestOnnxNumpy().test_zeros_none()
     unittest.main(verbosity=2)
diff --git a/_unittests/ut_npx/test_npx.py b/_unittests/ut_npx/test_npx.py
@@ -710,8 +710,8 @@ def impl(
         keys = list(sorted(f.onxs))
         self.assertIsInstance(f.onxs[keys[0]], ModelProto)
         k = keys[-1]
-        self.assertEqual(len(k), 3)
-        self.assertEqual(k[1:], ("axis", 0))
+        self.assertEqual(len(k), 4)
+        self.assertEqual(k[1:], ("axis", int, 0))
 
     def test_numpy_topk(self):
         f = topk(Input("X"), Input("K"))
@@ -2416,6 +2416,7 @@ def compute_labels(X, centers, use_sqrt=False):
             (DType(TensorProto.DOUBLE), 2),
             (DType(TensorProto.DOUBLE), 2),
             "use_sqrt",
+            bool,
             True,
         )
         self.assertEqual(f.available_versions, [key])
diff --git a/azure-pipelines.yml b/azure-pipelines.yml
@@ -48,7 +48,7 @@ jobs:
     vmImage: 'ubuntu-latest'
   strategy:
     matrix:
-      Python310-Linux:
+      Python311-Linux:
         python.version: '3.11'
     maxParallel: 3
 
@@ -96,7 +96,7 @@ jobs:
   strategy:
     matrix:
       Python310-Linux:
-        python.version: '3.11'
+        python.version: '3.10'
     maxParallel: 3
 
   steps:
@@ -149,7 +149,7 @@ jobs:
     vmImage: 'ubuntu-latest'
   strategy:
     matrix:
-      Python310-Linux:
+      Python311-Linux:
         python.version: '3.11'
     maxParallel: 3
 
@@ -202,7 +202,7 @@ jobs:
     vmImage: 'windows-latest'
   strategy:
     matrix:
-      Python310-Windows:
+      Python311-Windows:
         python.version: '3.11'
     maxParallel: 3
 
@@ -235,7 +235,7 @@ jobs:
     vmImage: 'macOS-latest'
   strategy:
     matrix:
-      Python310-Mac:
+      Python311-Mac:
         python.version: '3.11'
     maxParallel: 3
 
diff --git a/onnx_array_api/_helpers.py b/onnx_array_api/_helpers.py
@@ -39,7 +39,7 @@ def np_dtype_to_tensor_dtype(dtype: Any):
         elif dtype is int:
             dt = TensorProto.INT64
         elif dtype is float:
-            dt = TensorProto.FLOAT64
+            dt = TensorProto.DOUBLE
         else:
             raise KeyError(f"Unable to guess type for dtype={dtype}.")
     return dt
diff --git a/onnx_array_api/array_api/_onnx_common.py b/onnx_array_api/array_api/_onnx_common.py
@@ -44,7 +44,7 @@ def template_asarray(
             except OverflowError:
                 v = TEagerTensor(np.asarray(a, dtype=np.uint64))
     elif isinstance(a, float):
-        v = TEagerTensor(np.array(a, dtype=np.float32))
+        v = TEagerTensor(np.array(a, dtype=np.float64))
     elif isinstance(a, bool):
         v = TEagerTensor(np.array(a, dtype=np.bool_))
     elif isinstance(a, str):
diff --git a/onnx_array_api/array_api/onnx_numpy.py b/onnx_array_api/array_api/onnx_numpy.py
@@ -3,7 +3,6 @@
 """
 from typing import Any, Optional
 import numpy as np
-from onnx import TensorProto
 from ..npx.npx_functions import (
     all,
     abs,
@@ -16,10 +15,11 @@
     reshape,
     take,
 )
+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
+from ..npx.npx_types import DType, ElemType, TensorType, OptParType, ParType, Scalar
 from ._onnx_common import template_asarray
 from . import _finalize_array_api
 
@@ -31,6 +31,7 @@
     "astype",
     "empty",
     "equal",
+    "full",
     "isdtype",
     "isfinite",
     "isnan",
@@ -58,7 +59,7 @@ def asarray(
 
 def ones(
     shape: TensorType[ElemType.int64, "I", (None,)],
-    dtype: OptParType[DType] = DType(TensorProto.FLOAT),
+    dtype: OptParType[DType] = None,
     order: OptParType[str] = "C",
 ) -> TensorType[ElemType.numerics, "T"]:
     if isinstance(shape, tuple):
@@ -76,7 +77,7 @@ def ones(
 
 def empty(
     shape: TensorType[ElemType.int64, "I", (None,)],
-    dtype: OptParType[DType] = DType(TensorProto.FLOAT),
+    dtype: OptParType[DType] = None,
     order: OptParType[str] = "C",
 ) -> TensorType[ElemType.numerics, "T"]:
     raise RuntimeError(
@@ -87,7 +88,7 @@ def empty(
 
 def zeros(
     shape: TensorType[ElemType.int64, "I", (None,)],
-    dtype: OptParType[DType] = DType(TensorProto.FLOAT),
+    dtype: OptParType[DType] = None,
     order: OptParType[str] = "C",
 ) -> TensorType[ElemType.numerics, "T"]:
     if isinstance(shape, tuple):
@@ -103,6 +104,32 @@ def zeros(
     return generic_zeros(shape, dtype=dtype, order=order)
 
 
+def full(
+    shape: TensorType[ElemType.int64, "I", (None,)],
+    fill_value: ParType[Scalar] = None,
+    dtype: OptParType[DType] = None,
+    order: OptParType[str] = "C",
+) -> TensorType[ElemType.numerics, "T"]:
+    if fill_value is None:
+        raise TypeError("fill_value cannot be None")
+    value = fill_value
+    if isinstance(shape, tuple):
+        return generic_full(
+            EagerNumpyTensor(np.array(shape, dtype=np.int64)),
+            fill_value=value,
+            dtype=dtype,
+            order=order,
+        )
+    if isinstance(shape, int):
+        return generic_full(
+            EagerNumpyTensor(np.array([shape], dtype=np.int64)),
+            fill_value=value,
+            dtype=dtype,
+            order=order,
+        )
+    return generic_full(shape, fill_value=value, dtype=dtype, order=order)
+
+
 def _finalize():
     """
     Adds common attributes to Array API defined in this modules
diff --git a/onnx_array_api/npx/npx_core_api.py b/onnx_array_api/npx/npx_core_api.py
@@ -169,7 +169,7 @@ def wrapper(*inputs, **kwargs):
                 new_inputs.append(i)
             elif isinstance(i, (int, float)):
                 new_inputs.append(
-                    np.array([i], dtype=np.int64 if isinstance(i, int) else np.float32)
+                    np.array([i], dtype=np.int64 if isinstance(i, int) else np.float64)
                 )
             elif isinstance(i, str):
                 new_inputs.append(Input(i))
diff --git a/onnx_array_api/npx/npx_functions.py b/onnx_array_api/npx/npx_functions.py
@@ -15,14 +15,15 @@
     SequenceType,
     TensorType,
     TupleType,
+    Scalar,
 )
 from .npx_var import Var
 
 
 def _cstv(x):
     if isinstance(x, Var):
         return x
-    if isinstance(x, (int, float, np.ndarray)):
+    if isinstance(x, (int, float, bool, np.ndarray)):
         return cst(x)
     raise TypeError(f"Unexpected constant type {type(x)}.")
 
@@ -376,6 +377,42 @@ def expit(x: TensorType[ElemType.numerics, "T"]) -> TensorType[ElemType.numerics
     return var(x, op="Sigmoid")
 
 
+@npxapi_inline
+def full(
+    shape: TensorType[ElemType.int64, "I", (None,)],
+    dtype: OptParType[DType] = None,
+    fill_value: ParType[Scalar] = None,
+    order: OptParType[str] = "C",
+) -> TensorType[ElemType.numerics, "T"]:
+    """
+    Implements :func:`numpy.full`.
+    """
+    if order != "C":
+        raise RuntimeError(f"order={order!r} != 'C' not supported.")
+    if fill_value is None:
+        raise TypeError("fill_value cannot be None.")
+    if dtype is None:
+        if isinstance(fill_value, bool):
+            dtype = DType(TensorProto.BOOL)
+        elif isinstance(fill_value, int):
+            dtype = DType(TensorProto.INT64)
+        elif isinstance(fill_value, float):
+            dtype = DType(TensorProto.DOUBLE)
+        else:
+            raise TypeError(
+                f"Unexpected type {type(fill_value)} for fill_value={fill_value!r}."
+            )
+    if isinstance(fill_value, (float, int, bool)):
+        value = make_tensor(
+            name="cst", data_type=dtype.code, dims=[1], vals=[fill_value]
+        )
+    else:
+        raise NotImplementedError(
+            f"Unexpected type ({type(fill_value)} for fill_value={fill_value!r}."
+        )
+    return var(shape, value=value, op="ConstantOfShape")
+
+
 @npxapi_inline
 def floor(x: TensorType[ElemType.numerics, "T"]) -> TensorType[ElemType.numerics, "T"]:
     "See :func:`numpy.floor`."
@@ -464,7 +501,7 @@ def matmul(
 @npxapi_inline
 def ones(
     shape: TensorType[ElemType.int64, "I", (None,)],
-    dtype: OptParType[DType] = DType(TensorProto.FLOAT),
+    dtype: OptParType[DType] = None,
     order: OptParType[str] = "C",
 ) -> TensorType[ElemType.numerics, "T"]:
     """
@@ -473,7 +510,7 @@ def ones(
     if order != "C":
         raise RuntimeError(f"order={order!r} != 'C' not supported.")
     if dtype is None:
-        dtype = DType(TensorProto.FLOAT)
+        dtype = DType(TensorProto.DOUBLE)
     return var(
         shape,
         value=make_tensor(name="one", data_type=dtype.code, dims=[1], vals=[1]),
@@ -674,7 +711,7 @@ def where(
 @npxapi_inline
 def zeros(
     shape: TensorType[ElemType.int64, "I", (None,)],
-    dtype: OptParType[DType] = DType(TensorProto.FLOAT),
+    dtype: OptParType[DType] = None,
     order: OptParType[str] = "C",
 ) -> TensorType[ElemType.numerics, "T"]:
     """
@@ -683,7 +720,7 @@ def zeros(
     if order != "C":
         raise RuntimeError(f"order={order!r} != 'C' not supported.")
     if dtype is None:
-        dtype = DType(TensorProto.FLOAT)
+        dtype = DType(TensorProto.DOUBLE)
     return var(
         shape,
         value=make_tensor(name="zero", data_type=dtype.code, dims=[1], vals=[0]),
diff --git a/onnx_array_api/npx/npx_graph_builder.py b/onnx_array_api/npx/npx_graph_builder.py
@@ -784,7 +784,7 @@ def to_onnx(
                     node_inputs.append(input_name)
                     continue
 
-                if isinstance(i, (int, float)):
+                if isinstance(i, (int, float, bool)):
                     ni = np.array(i)
                     c = Cst(ni)
                     input_name = self._unique(var._prefix)
diff --git a/onnx_array_api/npx/npx_jit_eager.py b/onnx_array_api/npx/npx_jit_eager.py
diff --git a/onnx_array_api/npx/npx_numpy_tensors_ops.py b/onnx_array_api/npx/npx_numpy_tensors_ops.py
diff --git a/onnx_array_api/npx/npx_tensors.py b/onnx_array_api/npx/npx_tensors.py
diff --git a/onnx_array_api/npx/npx_types.py b/onnx_array_api/npx/npx_types.py
diff --git a/onnx_array_api/npx/npx_var.py b/onnx_array_api/npx/npx_var.py

Original file line number	Diff line number	Diff line change
`@@ -169,7 +169,7 @@ def wrapper(inputs, *kwargs):`
`169`	`169`	`new_inputs.append(i)`
`170`	`170`	`elif isinstance(i, (int, float)):`
`171`	`171`	`new_inputs.append(`
`172`		`- np.array([i], dtype=np.int64 if isinstance(i, int) else np.float32)`
	`172`	`+ np.array([i], dtype=np.int64 if isinstance(i, int) else np.float64)`
`173`	`173`	`)`
`174`	`174`	`elif isinstance(i, str):`
`175`	`175`	`new_inputs.append(Input(i))`