Merge pull request #2 from sdpython/array_api

Implements the first step for the Array API, test with scikit-learn

Author: xadupre

_unittests/ut__main/test_documentation_examples.py

@@ -4,8 +4,12 @@
 import importlib
 import subprocess
 import time
+from onnx_array_api import __file__ as onnx_array_api_file
 from onnx_array_api.ext_test_case import ExtTestCase
 
+VERBOSE = 0
+ROOT = os.path.realpath(os.path.abspath(os.path.join(onnx_array_api_file, "..", "..")))
+
 
 def import_source(module_file_path, module_name):
     if not os.path.exists(module_file_path):
@@ -20,43 +24,58 @@ def import_source(module_file_path, module_name):
 
 
 class TestDocumentationExamples(ExtTestCase):
-    def test_documentation_examples(self):
+    def run_test(self, fold: str, name: str, verbose=0) -> int:
+        ppath = os.environ.get("PYTHONPATH", "")
+        if len(ppath) == 0:
+            os.environ["PYTHONPATH"] = ROOT
+        elif ROOT not in ppath:
+            sep = ";" if sys.platform == "win32" else ":"
+            os.environ["PYTHONPATH"] = ppath + sep + ROOT
+        perf = time.perf_counter()
+        try:
+            mod = import_source(fold, os.path.splitext(name)[0])
+            assert mod is not None
+        except FileNotFoundError:
+            # try another way
+            cmds = [sys.executable, "-u", os.path.join(fold, name)]
+            p = subprocess.Popen(cmds, stdout=subprocess.PIPE, stderr=subprocess.PIPE)
+            res = p.communicate()
+            out, err = res
+            st = err.decode("ascii", errors="ignore")
+            if len(st) > 0 and "Traceback" in st:
+                if '"dot" not found in path.' in st:
+                    # dot not installed, this part
+                    # is tested in onnx framework
+                    if verbose:
+                        print(f"failed: {name!r} due to missing dot.")
+                    return 0
+                raise AssertionError(
+                    "Example '{}' (cmd: {} - exec_prefix='{}') "
+                    "failed due to\n{}"
+                    "".format(name, cmds, sys.exec_prefix, st)
+                )
+        dt = time.perf_counter() - perf
+        if verbose:
+            print(f"{dt:.3f}: run {name!r}")
+        return 1
+
+    @classmethod
+    def add_test_methods(cls):
         this = os.path.abspath(os.path.dirname(__file__))
         fold = os.path.normpath(os.path.join(this, "..", "..", "_doc", "examples"))
         found = os.listdir(fold)
-        tested = 0
         for name in found:
             if name.startswith("plot_") and name.endswith(".py"):
-                perf = time.perf_counter()
-                try:
-                    mod = import_source(fold, os.path.splitext(name)[0])
-                    assert mod is not None
-                except FileNotFoundError:
-                    # try another way
-                    cmds = [sys.executable, "-u", os.path.join(fold, name)]
-                    p = subprocess.Popen(
-                        cmds, stdout=subprocess.PIPE, stderr=subprocess.PIPE
-                    )
-                    res = p.communicate()
-                    out, err = res
-                    st = err.decode("ascii", errors="ignore")
-                    if len(st) > 0 and "Traceback" in st:
-                        if '"dot" not found in path.' in st:
-                            # dot not installed, this part
-                            # is tested in onnx framework
-                            print(f"failed: {name!r} due to missing dot.")
-                            continue
-                        raise AssertionError(
-                            "Example '{}' (cmd: {} - exec_prefix='{}') "
-                            "failed due to\n{}"
-                            "".format(name, cmds, sys.exec_prefix, st)
-                        )
-                dt = time.perf_counter() - perf
-                print(f"{dt:.3f}: run {name!r}")
-                tested += 1
-        if tested == 0:
-            raise AssertionError("No example was tested.")
+                short_name = os.path.split(os.path.splitext(name)[0])[-1]
+
+                def _test_(self, name=name):
+                    res = self.run_test(fold, name, verbose=VERBOSE)
+                    self.assertTrue(res)
+
+                setattr(cls, f"test_{short_name}", _test_)
+
 
+TestDocumentationExamples.add_test_methods()
 
 if __name__ == "__main__":
-    unittest.main()
+    unittest.main(verbosity=2)

_unittests/ut_npx/test_npx.py

@@ -71,6 +71,7 @@
 from onnx_array_api.npx.npx_functions import sinh as sinh_inline
 from onnx_array_api.npx.npx_functions import sqrt as sqrt_inline
 from onnx_array_api.npx.npx_functions import squeeze as squeeze_inline
+from onnx_array_api.npx.npx_functions import take as take_inline
 from onnx_array_api.npx.npx_functions import tan as tan_inline
 from onnx_array_api.npx.npx_functions import tanh as tanh_inline
 from onnx_array_api.npx.npx_functions import topk as topk_inline
@@ -2439,7 +2440,19 @@ def myloss(x, y):
         res_eager = eager_myloss(x, y)
         assert_allclose(res_jit, res_eager)
 
+    def test_take(self):
+        data = np.random.randn(3, 3).astype(np.float32)
+        indices = np.array([[0, 2]])
+        y = np.take(data, indices, axis=1)
+
+        f = take_inline(Input("A"), Input("B"), axis=1)
+        self.assertIsInstance(f, Var)
+        onx = f.to_onnx(constraints={"A": Float64[None], "B": Int64[None]})
+        ref = ReferenceEvaluator(onx)
+        got = ref.run(None, {"A": data, "B": indices})
+        self.assertEqualArray(y, got[0])
+
 
 if __name__ == "__main__":
-    TestNpx().test_eager_cst_index()
+    TestNpx().test_take()
     unittest.main(verbosity=2)

_unittests/ut_npx/test_sklearn_array_api.py

@@ -0,0 +1,43 @@
+import unittest
+import numpy as np
+from onnx.defs import onnx_opset_version
+from sklearn import config_context
+from sklearn.discriminant_analysis import LinearDiscriminantAnalysis
+from onnx_array_api.ext_test_case import ExtTestCase
+from onnx_array_api.npx.npx_numpy_tensors import EagerNumpyTensor
+
+
+DEFAULT_OPSET = onnx_opset_version()
+
+
+def take(self, X, indices, *, axis):
+    # Overwritting method take as it is using iterators.
+    # When array_api supports `take` we can use this directly
+    # https://github.com/data-apis/array-api/issues/177
+    X_np = self._namespace.take(X, indices, axis=axis)
+    return self._namespace.asarray(X_np)
+
+
+class TestSklearnArrayAPI(ExtTestCase):
+    def test_sklearn_array_api_linear_discriminant(self):
+        from sklearn.utils._array_api import _ArrayAPIWrapper
+
+        _ArrayAPIWrapper.take = take
+        X = np.array([[-1, -1], [-2, -1], [-3, -2], [1, 1], [2, 1], [3, 2]])
+        y = np.array([1, 1, 1, 2, 2, 2])
+        ana = LinearDiscriminantAnalysis()
+        ana = LinearDiscriminantAnalysis()
+        ana.fit(X, y)
+        expected = ana.predict(X)
+
+        new_x = EagerNumpyTensor(X)
+        self.assertStartsWith("EagerNumpyTensor(array([[", repr(new_x))
+        with config_context(array_api_dispatch=True):
+            got = ana.predict(new_x)
+        self.assertEqualArray(expected, got.numpy())
+
+
+if __name__ == "__main__":
+    # import logging
+    # logging.basicConfig(level=logging.DEBUG)
+    unittest.main(verbosity=2)

onnx_array_api/__init__.py

@@ -3,5 +3,5 @@
 (Numpy) Array API for ONNX.
 """
 
-__version__ = "0.1.1"
+__version__ = "0.1.2"
 __author__ = "Xavier Dupré"

onnx_array_api/npx/npx_array_api.py

@@ -10,6 +10,14 @@ class ArrayApi:
     List of supported method by a tensor.
     """
 
+    def __array_namespace__(self):
+        """
+        Returns the module holding all the available functions.
+        """
+        from onnx_array_api.npx import npx_functions
+
+        return npx_functions
+
     def generic_method(self, method_name, *args: Any, **kwargs: Any) -> Any:
         raise NotImplementedError(
             f"Method {method_name!r} must be overwritten "

onnx_array_api/npx/npx_core_api.py

@@ -123,6 +123,7 @@ def wrapper(*inputs, **kwargs):
             for x in inputs:
                 if isinstance(x, EagerTensor):
                     tensor_class = x.__class__
+                    break
             if tensor_class is None:
                 raise RuntimeError(
                     f"Unable to find an EagerTensor in types "

onnx_array_api/npx/npx_functions.py

@@ -1,11 +1,13 @@
-from typing import Optional, Tuple, Union
+from typing import Any, Optional, Tuple, Union
 
 import numpy as np
-from onnx import FunctionProto, ModelProto, NodeProto
+from onnx import FunctionProto, ModelProto, NodeProto, TensorProto
+from onnx.helper import np_dtype_to_tensor_dtype
 from onnx.numpy_helper import from_array
 
 from .npx_constants import FUNCTION_DOMAIN
 from .npx_core_api import cst, make_tuple, npxapi_inline, var
+from .npx_tensors import ArrayApi
 from .npx_types import (
     ElemType,
     OptParType,
@@ -155,6 +157,50 @@ 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, ArrayApi):
+        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[int] = 1
+) -> TensorType[ElemType.numerics, "T2"]:
+    """
+    Cast an array.
+    """
+    if isinstance(dtype, Var):
+        raise TypeError(
+            f"dtype is an attribute, it cannot be a Variable of type {type(dtype)}."
+        )
+    try:
+        to = np_dtype_to_tensor_dtype(dtype)
+    except KeyError:
+        if dtype is int:
+            to = TensorProto.INT64
+        elif dtype is float:
+            to = TensorProto.float64
+        else:
+            raise ValueError(f"Unable to guess tensor type from {dtype}.")
+    return var(a, op="Cast", to=to)
+
+
 @npxapi_inline
 def cdist(
     xa: TensorType[ElemType.numerics, "T"],
@@ -412,6 +458,17 @@ def relu(x: TensorType[ElemType.numerics, "T"]) -> TensorType[ElemType.numerics,
     return var(x, op="Relu")
 
 
+@npxapi_inline
+def reshape(
+    x: TensorType[ElemType.numerics, "T"], shape: TensorType[ElemType.int64, "I"]
+) -> TensorType[ElemType.numerics, "T"]:
+    "See :func:`numpy.reshape`."
+    if isinstance(shape, int):
+        shape = cst(np.array([shape], dtype=np.int64))
+    shape_reshaped = var(shape, cst(np.array([-1], dtype=np.int64)), op="Reshape")
+    return var(x, shape_reshaped, op="Reshape")
+
+
 @npxapi_inline
 def round(x: TensorType[ElemType.numerics, "T"]) -> TensorType[ElemType.numerics, "T"]:
     "See :func:`numpy.round`."
@@ -469,6 +526,16 @@ def squeeze(
     return var(x, axis, op="Squeeze")
 
 
+@npxapi_inline
+def take(
+    data: TensorType[ElemType.numerics, "T"],
+    indices: TensorType[ElemType.int64, "I"],
+    axis: ParType[int] = 0,
+) -> TensorType[ElemType.numerics, "T"]:
+    "See :func:`numpy.take`."
+    return var(data, indices, op="Gather", axis=axis)
+
+
 @npxapi_inline
 def tan(x: TensorType[ElemType.numerics, "T"]) -> TensorType[ElemType.numerics, "T"]:
     "See :func:`numpy.tan`."

onnx_array_api/npx/npx_graph_builder.py

@@ -235,7 +235,13 @@ def make_node(
                 f"Cannot create a node Identity for {len(inputs)} input(s) and "
                 f"{len(outputs)} output(s)."
             )
-        node = make_node(op, inputs, outputs, domain=domain, **new_kwargs)
+        try:
+            node = make_node(op, inputs, outputs, domain=domain, **new_kwargs)
+        except TypeError as e:
+            raise TypeError(
+                f"Unable to create node {op!r}, with inputs={inputs}, "
+                f"outputs={outputs}, domain={domain!r}, new_kwargs={new_kwargs}."
+            ) from e
         for p in protos:
             node.attribute.append(p)
         if attribute_protos is not None: