improvments

sdpython · sdpython · Jun 10, 2023 · Jun 5, 2023 · Jun 5, 2023 · Jun 5, 2023
commit 1e218b5a3e84d3db8b44f8ef3b1db248a660bc21
diff --git a/_unittests/test_array_api.sh b/_unittests/test_array_api.sh
@@ -0,0 +1,2 @@
+export ARRAY_API_TESTS_MODULE=onnx_array_api.array_api.onnx_numpy
+pytest ../array-api-tests/array_api_tests/test_creation_functions.py::test_zeros
diff --git a/_unittests/ut_npx/test_npx.py b/_unittests/ut_npx/test_npx.py
@@ -29,6 +29,7 @@
     npxapi_inline,
 )
 from onnx_array_api.npx.npx_functions import absolute as absolute_inline
+from onnx_array_api.npx.npx_functions import all as all_inline
 from onnx_array_api.npx.npx_functions import arange as arange_inline
 from onnx_array_api.npx.npx_functions import arccos as arccos_inline
 from onnx_array_api.npx.npx_functions import arccosh as arccosh_inline
@@ -50,6 +51,7 @@
 from onnx_array_api.npx.npx_functions import det as det_inline
 from onnx_array_api.npx.npx_functions import dot as dot_inline
 from onnx_array_api.npx.npx_functions import einsum as einsum_inline
+from onnx_array_api.npx.npx_functions import equal as equal_inline
 from onnx_array_api.npx.npx_functions import erf as erf_inline
 from onnx_array_api.npx.npx_functions import exp as exp_inline
 from onnx_array_api.npx.npx_functions import expand_dims as expand_dims_inline
@@ -1163,6 +1165,16 @@ def test_astype(self):
         got = ref.run(None, {"A": x})
         self.assertEqualArray(z, got[0])
 
+    def test_astype_dtype(self):
+        f = absolute_inline(copy_inline(Input("A")).astype(DType(7)))
+        self.assertIsInstance(f, Var)
+        onx = f.to_onnx(constraints={"A": Float64[None]})
+        x = np.array([[-5.4, 6.6]], dtype=np.float64)
+        z = np.abs(x.astype(np.int64))
+        ref = ReferenceEvaluator(onx)
+        got = ref.run(None, {"A": x})
+        self.assertEqualArray(z, got[0])
+
     def test_astype_int(self):
         f = absolute_inline(copy_inline(Input("A")).astype(1))
         self.assertIsInstance(f, Var)
@@ -1421,6 +1433,9 @@ def test_einsum(self):
             lambda x, y: np.einsum(equation, x, y),
         )
 
+    def test_equal(self):
+        self.common_test_inline_bin(equal_inline, np.equal)
+
     @unittest.skipIf(scipy is None, reason="scipy is not installed.")
     def test_erf(self):
         self.common_test_inline(erf_inline, scipy.special.erf)
@@ -1468,7 +1483,8 @@ def test_hstack(self):
     def test_identity(self):
         f = identity_inline(2, dtype=np.float64)
         onx = f.to_onnx(constraints={(0, False): Float64[None]})
-        z = np.identity(2)
+        self.assertIn("dtype:", str(onx))
+        z = np.identity(2).astype(np.float64)
         ref = ReferenceEvaluator(onx)
         got = ref.run(None, {})
         self.assertEqualArray(z, got[0])
@@ -2465,7 +2481,18 @@ def test_take(self):
         got = ref.run(None, {"A": data, "B": indices})
         self.assertEqualArray(y, got[0])
 
+    def test_numpy_all(self):
+        data = np.array([[1, 0], [1, 1]]).astype(np.bool_)
+        y = np.all(data, axis=1)
+
+        f = all_inline(Input("A"), axis=1)
+        self.assertIsInstance(f, Var)
+        onx = f.to_onnx(constraints={"A": Bool[None]})
+        ref = ReferenceEvaluator(onx)
+        got = ref.run(None, {"A": data})
+        self.assertEqualArray(y, got[0])
+
 
 if __name__ == "__main__":
-    TestNpx().test_take()
+    TestNpx().test_identity()
     unittest.main(verbosity=2)
diff --git a/azure-pipelines.yml b/azure-pipelines.yml
@@ -125,14 +125,12 @@ jobs:
   - script: |
       export ARRAY_API_TESTS_MODULE=onnx_array_api.array_api.onnx_numpy
       cd array-api-tests
+    displayName: 'Set API'
+  - script: |
       python -m pytest -xv array_api_tests/test_creation_functions.py::test_zeros
-      cd ..
     displayName: "test_creation_functions.py::test_zeros"
   - script: |
-      export ARRAY_API_TESTS_MODULE=onnx_array_api.array_api.onnx_numpy
-      cd array-api-tests
       python -m pytest -x array_api_tests
-      cd ..
     displayName: "all tests"
 
 - job: 'TestLinux'

diff --git a/onnx_array_api/array_api/onnx_numpy.py b/onnx_array_api/array_api/onnx_numpy.py
@@ -6,14 +6,16 @@
 from onnx import TensorProto
 from ..npx.npx_array_api import BaseArrayApi
 from ..npx.npx_functions import (
+    all,
     abs,
     absolute,
     astype,
+    copy as copy_inline,
+    equal,
     isdtype,
     reshape,
     take,
 )
-from ..npx.npx_functions import asarray as generic_asarray
 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
@@ -22,8 +24,10 @@
 __all__ = [
     "abs",
     "absolute",
+    "all",
     "asarray",
     "astype",
+    "equal",
     "isdtype",
     "reshape",
     "take",
@@ -41,13 +45,36 @@ def asarray(
     """
     Converts anything into an array.
     """
+    if order not in ("C", None):
+        raise NotImplementedError(f"asarray is not implemented for order={order!r}.")
+    if like is not None:
+        raise NotImplementedError(
+            f"asarray is not implemented for like != None (type={type(like)})."
+        )
     if isinstance(a, BaseArrayApi):
-        return generic_asarray(a, dtype=dtype, order=order, like=like, copy=copy)
+        if copy:
+            if dtype is None:
+                return copy_inline(a)
+            return copy_inline(a).astype(dtype=dtype)
+        if dtype is None:
+            return a
+        return a.astype(dtype=dtype)
+
     if isinstance(a, int):
-        return EagerNumpyTensor(np.array(a, dtype=np.int64))
-    if isinstance(a, float):
-        return EagerNumpyTensor(np.array(a, dtype=np.float32))
-    raise NotImplementedError(f"asarray not implemented for type {type(a)}.")
+        v = EagerNumpyTensor(np.array(a, dtype=np.int64))
+    elif isinstance(a, float):
+        v = EagerNumpyTensor(np.array(a, dtype=np.float32))
+    elif isinstance(a, bool):
+        v = EagerNumpyTensor(np.array(a, dtype=np.bool_))
+    elif isinstance(a, str):
+        v = EagerNumpyTensor(np.array(a, dtype=np.str_))
+    else:
+        raise RuntimeError(f"Unexpected type {type(a)} for the first input.")
+    if dtype is not None:
+        vt = v.astype(dtype=dtype)
+    else:
+        vt = v
+    return vt
 
 
 def zeros(

diff --git a/onnx_array_api/array_api/onnx_ort.py b/onnx_array_api/array_api/onnx_ort.py
@@ -1,51 +1,30 @@
 """
 Array API valid for an :class:`EagerOrtTensor`.
 """
-from typing import Any, Optional
-import numpy as np
-from ..npx.npx_array_api import BaseArrayApi
 from ..npx.npx_functions import (
+    all,
     abs,
     absolute,
     astype,
+    equal,
     isdtype,
     reshape,
     take,
 )
-from ..npx.npx_functions import asarray as generic_asarray
-from ..ort.ort_tensors import EagerOrtTensor
 from . import _finalize_array_api
 
 __all__ = [
+    "all",
     "abs",
     "absolute",
-    "asarray",
     "astype",
+    "equal",
     "isdtype",
     "reshape",
     "take",
 ]
 
 
-def asarray(
-    a: Any,
-    dtype: Any = None,
-    order: Optional[str] = None,
-    like: Any = None,
-    copy: bool = False,
-):
-    """
-    Converts anything into an array.
-    """
-    if isinstance(a, BaseArrayApi):
-        return generic_asarray(a, dtype=dtype, order=order, like=like, copy=copy)
-    if isinstance(a, int):
-        return EagerOrtTensor(np.array(a, dtype=np.int64))
-    if isinstance(a, float):
-        return EagerOrtTensor(np.array(a, dtype=np.float32))
-    raise NotImplementedError(f"asarray not implemented for type {type(a)}.")
-
-
 def _finalize():
     from . import onnx_ort
 

diff --git a/onnx_array_api/npx/npx_functions.py b/onnx_array_api/npx/npx_functions.py
@@ -1,4 +1,4 @@
-from typing import Any, Optional, Tuple, Union
+from typing import Optional, Tuple, Union
 
 import array_api_compat.numpy as np_array_api
 import numpy as np
@@ -8,7 +8,6 @@
 
 from .npx_constants import FUNCTION_DOMAIN
 from .npx_core_api import cst, make_tuple, npxapi_inline, npxapi_no_inline, var
-from .npx_tensors import BaseArrayApi
 from .npx_types import (
     DType,
     ElemType,
@@ -43,6 +42,27 @@ def absolute(
     return var(x, op="Abs")
 
 
+@npxapi_inline
+def all(
+    x: TensorType[ElemType.bool_, "T"],
+    axis: Optional[TensorType[ElemType.int64, "I"]] = None,
+    keepdims: ParType[int] = 0,
+) -> TensorType[ElemType.bool_, "T"]:
+    "See :func:`numpy.all`."
+
+    xi = var(x, op="Cast", to=TensorProto.INT64)
+
+    if axis is None:
+        red = xi.min(keepdims=keepdims)
+    else:
+        if isinstance(axis, int):
+            axis = [axis]
+        if isinstance(axis, (tuple, list)):
+            axis = cst(np.array(axis, dtype=np.int64))
+        red = xi.min(axis, keepdims=keepdims)
+    return var(red, cst(1), op="Equal")
+
+
 @npxapi_inline
 def arccos(x: TensorType[ElemType.numerics, "T"]) -> TensorType[ElemType.numerics, "T"]:
     "See :func:`numpy.arccos`."
@@ -159,27 +179,6 @@ def arctanh(
     return var(x, op="Atanh")
 
 
-def asarray(
-    a: Any,
-    dtype: Any = None,
-    order: Optional[str] = None,
-    like: Any = None,
-    copy: bool = False,
-):
-    """
-    Converts anything into an array.
-    """
-    if dtype is not None:
-        raise RuntimeError("Method 'astype' should be used to change the type.")
-    if order is not None:
-        raise NotImplementedError(f"order={order!r} not implemented.")
-    if isinstance(a, BaseArrayApi):
-        if copy:
-            return a.__class__(a, copy=copy)
-        return a
-    raise NotImplementedError(f"asarray not implemented for type {type(a)}.")
-
-
 @npxapi_inline
 def astype(
     a: TensorType[ElemType.numerics, "T1"], dtype: OptParType[DType] = 1
@@ -335,6 +334,14 @@ def einsum(
     return var(*x, op="Einsum", equation=equation)
 
 
+@npxapi_inline
+def equal(
+    x: TensorType[ElemType.allowed, "T"], y: TensorType[ElemType.allowed, "T"]
+) -> TensorType[ElemType.bool_, "T1"]:
+    "See :func:`numpy.isnan`."
+    return var(x, y, op="Equal")
+
+
 @npxapi_inline
 def erf(x: TensorType[ElemType.numerics, "T"]) -> TensorType[ElemType.numerics, "T"]:
     "See :func:`scipy.special.erf`."
@@ -382,18 +389,20 @@ def hstack(
 
 
 @npxapi_inline
-def copy(x: TensorType[ElemType.numerics, "T"]) -> TensorType[ElemType.numerics, "T"]:
+def copy(x: TensorType[ElemType.allowed, "T"]) -> TensorType[ElemType.allowed, "T"]:
     "Makes a copy."
     return var(x, op="Identity")
 
 
 @npxapi_inline
-def identity(n: ParType[int], dtype=None) -> TensorType[ElemType.numerics, "T"]:
+def identity(
+    n: ParType[int], dtype: OptParType[DType] = None
+) -> TensorType[ElemType.numerics, "T"]:
     "Makes a copy."
-    val = np.array([n, n], dtype=np.int64)
-    shape = cst(val)
     model = var(
-        shape, op="ConstantOfShape", value=from_array(np.array([0], dtype=np.int64))
+        cst(np.array([n, n], dtype=np.int64)),
+        op="ConstantOfShape",
+        value=from_array(np.array([0], dtype=np.int64)),
     )
     v = var(model, dtype=dtype, op="EyeLike")
     return v
@@ -416,7 +425,7 @@ def isdtype(
 
 
 @npxapi_inline
-def isnan(x: TensorType[ElemType.numerics, "T"]) -> TensorType[ElemType.bool_, "T"]:
+def isnan(x: TensorType[ElemType.numerics, "T"]) -> TensorType[ElemType.bool_, "T1"]:
     "See :func:`numpy.isnan`."
     return var(x, op="IsNaN")
 
@@ -643,6 +652,8 @@ def zeros(
     """
     if order != "C":
         raise RuntimeError(f"order={order!r} != 'C' not supported.")
+    if dtype is None:
+        dtype = DType(TensorProto.FLOAT)
     return var(
         shape,
         value=make_tensor(name="zero", data_type=dtype.code, dims=[1], vals=[0]),

diff --git a/onnx_array_api/npx/npx_jit_eager.py b/onnx_array_api/npx/npx_jit_eager.py
@@ -168,6 +168,8 @@ def make_key(*values, **kwargs):
                         else:
                             newv.append(t)
                     res.append(tuple(newv))
+                elif v is None and k in {"dtype"}:
+                    continue
                 else:
                     raise TypeError(
                         f"Type {type(v)} is not yet supported, "
@@ -520,6 +522,8 @@ def _preprocess_constants(self, *args):
             elif isinstance(n, (int, float)):
                 new_args.append(self.tensor_class(np.array(n)))
                 modified = True
+            elif isinstance(n, DType):
+                new_args.append(n)
             elif n in (int, float):
                 # usually used to cast
                 new_args.append(n)
@@ -554,7 +558,17 @@ def __call__(self, *args, already_eager=False, **kwargs):
                     lambda t: t is not None
                     and not isinstance(
                         t,
-                        (EagerTensor, Cst, int, float, tuple, slice, type, np.ndarray),
+                        (
+                            EagerTensor,
+                            Cst,
+                            int,
+                            float,
+                            tuple,
+                            slice,
+                            type,
+                            np.ndarray,
+                            DType,
+                        ),
                     ),
                     args,
                 )
Original file line number	Diff line number	Diff line change
		@@ -0,0 +1,2 @@
		export ARRAY_API_TESTS_MODULE=onnx_array_api.array_api.onnx_numpy
		pytest ../array-api-tests/array_api_tests/test_creation_functions.py::test_zeros