first step for the array api

sdpython · sdpython · commit dd1876f3349c · 2023-03-19T13:33:14.000+01:00
diff --git a/_unittests/ut_npx/test_sklearn_array_api.py b/_unittests/ut_npx/test_sklearn_array_api.py
@@ -0,0 +1,30 @@
+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()
+
+
+class TestSklearnArrayAPI(ExtTestCase):
+    def test_sklearn_array_api_linear_discriminant(self):
+        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)
+
+
+if __name__ == "__main__":
+    unittest.main(verbosity=2)
diff --git a/onnx_array_api/npx/npx_array_api.py b/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 "
diff --git a/onnx_array_api/npx/npx_functions.py b/onnx_array_api/npx/npx_functions.py
@@ -1,11 +1,12 @@
-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.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 +156,38 @@ 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.
+    """
+    g = a.astype(dtype)
+    return g
+
+
 @npxapi_inline
 def cdist(
     xa: TensorType[ElemType.numerics, "T"],
@@ -412,6 +445,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`."
diff --git a/onnx_array_api/npx/npx_jit_eager.py b/onnx_array_api/npx/npx_jit_eager.py
@@ -338,6 +338,9 @@ def _preprocess_constants(self, *args):
             elif isinstance(n, (int, float)):
                 new_args.append(self.tensor_class(np.array(n)))
                 modified = True
+            elif n in (int, float):
+                # usually used to cast
+                new_args.append(n)
             elif n is None:
                 new_args.append(n)
             else:
@@ -365,7 +368,9 @@ def __call__(self, *args, already_eager=False, **kwargs):
             if any(
                 map(
                     lambda t: t is not None
-                    and not isinstance(t, (EagerTensor, Cst, int, float, tuple, slice)),
+                    and not isinstance(
+                        t, (EagerTensor, Cst, int, float, tuple, slice, type)
+                    ),
                     args,
                 )
             ):
diff --git a/onnx_array_api/npx/npx_numpy_tensors.py b/onnx_array_api/npx/npx_numpy_tensors.py
@@ -70,6 +70,10 @@ def __init__(self, tensor: np.ndarray):
         else:
             raise TypeError(f"A numpy array is expected not {type(tensor)}.")
 
+    def __repr__(self) -> str:
+        "usual"
+        return f"{self.__class__.__name__}({self._tensor!r})"
+
     def numpy(self):
         "Returns the array converted into a numpy array."
         return self._tensor
@@ -107,6 +111,11 @@ def dims(self):
             return self._tensor.shape
         return (None,) + self._tensor.shape[1:]
 
+    @property
+    def ndim(self):
+        "Returns the number of dimensions (rank)."
+        return len(self.shape)
+
     @property
     def shape(self) -> Tuple[int, ...]:
         "Returns the shape of the tensor."
diff --git a/onnx_array_api/npx/npx_tensors.py b/onnx_array_api/npx/npx_tensors.py
@@ -1,5 +1,8 @@
 from typing import Any
 
+import numpy as np
+from onnx.helper import np_dtype_to_tensor_dtype
+
 from .npx_array_api import ArrayApi
 
 
@@ -59,6 +62,16 @@ def _getitem_impl_var(obj, index, method_name=None):
         meth = getattr(Var, method_name)
         return meth(obj, index)
 
+    @staticmethod
+    def _astype_impl(x, dtype, method_name=None):
+        # avoids circular imports.
+        from .npx_var import Var
+
+        if not isinstance(x, Var):
+            raise TypeError(f"Input 0 must be a Var not {type(x)}.")
+        meth = getattr(Var, "astype")
+        return meth(x, dtype)
+
     @staticmethod
     def _getitem_impl_tuple(obj, index=None, method_name=None):
         # avoids circular imports.
@@ -69,13 +82,114 @@ def _getitem_impl_tuple(obj, index=None, method_name=None):
         meth = getattr(Var, method_name)
         return meth(obj, index)
 
+    def _generic_method_getitem(self, method_name, *args: Any, **kwargs: Any) -> Any:
+        # avoids circular imports.
+        from .npx_jit_eager import eager_onnx
+
+        if len(args) != 1:
+            raise ValueError(
+                f"Unexpected number of argument {len(args)}, it should be one."
+            )
+        if isinstance(args[0], tuple):
+            eag = eager_onnx(
+                EagerTensor._getitem_impl_tuple, self.__class__, bypass_eager=True
+            )
+            res = eag(self, index=args[0], method_name=method_name, already_eager=True)
+        else:
+            eag = eager_onnx(
+                EagerTensor._getitem_impl_var, self.__class__, bypass_eager=True
+            )
+            res = eag(self, args[0], method_name=method_name, already_eager=True)
+        if isinstance(res, tuple) and len(res) == 1:
+            return res[0]
+        return res
+
+    def _generic_method_operator(self, method_name, *args: Any, **kwargs: Any) -> Any:
+        # avoids circular imports.
+        from .npx_jit_eager import eager_onnx
+
+        if len(args) not in (0, 1):
+            raise ValueError(
+                f"An operator must have zero or one argument not {len(args)}."
+            )
+        if len(kwargs) not in (0, 1):
+            raise ValueError(f"Operators do not support parameters {len(kwargs)}.")
+
+        # let's cast numpy arrays into constants.
+        new_args = []
+        for a in args:
+            if isinstance(a, np.ndarray):
+                new_args.append(self.__class__(a).astype(self.dtype))
+            else:
+                new_args.append(a)
+
+        eag = eager_onnx(EagerTensor._op_impl, self.__class__, bypass_eager=True)
+        res = eag(self, *new_args, method_name=method_name, already_eager=True)
+        if isinstance(res, tuple) and len(res) == 1:
+            return res[0]
+        return res
+
+    def _generic_method_reduce(self, method_name, *args: Any, **kwargs: Any) -> Any:
+        # avoids circular imports.
+        from .npx_jit_eager import eager_onnx
+
+        if len(args) not in (0, 1):
+            raise ValueError(
+                f"An operator must have zero or one argument not {len(args)}."
+            )
+
+        if "axis" in kwargs:
+            axes = kwargs["axis"]
+            del kwargs["axis"]
+        else:
+            axes = None
+        if axes is None:
+            eag = eager_onnx(
+                EagerTensor._reduce_impl_noaxes, self.__class__, bypass_eager=True
+            )
+            res = eag(self, method_name=method_name, already_eager=True, **kwargs)
+        else:
+            eag = eager_onnx(
+                EagerTensor._reduce_impl, self.__class__, bypass_eager=True
+            )
+            res = eag(self, axes, method_name=method_name, already_eager=True, **kwargs)
+        if isinstance(res, tuple) and len(res) == 1:
+            return res[0]
+        return res
+
+    @staticmethod
+    def _np_dtype_to_tensor_dtype(dtype):
+        if dtype == int:
+            dtype = np.dtype("int64")
+        elif dtype == float:
+            dtype = np.dtype("float64")
+        return np_dtype_to_tensor_dtype(dtype)
+
+    def _generic_method_astype(self, method_name, *args: Any, **kwargs: Any) -> Any:
+        # avoids circular imports.
+        from .npx_jit_eager import eager_onnx
+        from .npx_var import Var
+
+        if len(args) != 1:
+            raise ValueError(f"astype takes only one argument not {len(args)}.")
+
+        dtype = (
+            args[0]
+            if isinstance(args[0], (int, Var))
+            else self._np_dtype_to_tensor_dtype(args[0])
+        )
+        eag = eager_onnx(EagerTensor._astype_impl, self.__class__, bypass_eager=True)
+        res = eag(self, dtype, method_name=method_name, already_eager=True, **kwargs)
+        if isinstance(res, tuple) and len(res) == 1:
+            return res[0]
+        return res
+
     def generic_method(self, method_name, *args: Any, **kwargs: Any) -> Any:
         """
         The method converts the method into an ONNX graph build by the
         corresponding method in class Var.
         """
         # avoids circular imports.
-        from .npx_jit_eager import eager_onnx
         from .npx_var import Var
 
         if not hasattr(Var, method_name):
@@ -84,70 +198,18 @@ def generic_method(self, method_name, *args: Any, **kwargs: Any) -> Any:
                 f"This method cannot be converted into an ONNX graph."
             )
         if method_name == "__getitem__":
-            if len(args) != 1:
-                raise ValueError(
-                    f"Unexpected number of argument {len(args)}, it should be one."
-                )
-            if isinstance(args[0], tuple):
-                eag = eager_onnx(
-                    EagerTensor._getitem_impl_tuple, self.__class__, bypass_eager=True
-                )
-                res = eag(
-                    self, index=args[0], method_name=method_name, already_eager=True
-                )
-            else:
-                eag = eager_onnx(
-                    EagerTensor._getitem_impl_var, self.__class__, bypass_eager=True
-                )
-                res = eag(self, args[0], method_name=method_name, already_eager=True)
-            if isinstance(res, tuple) and len(res) == 1:
-                return res[0]
-            return res
+            return self._generic_method_getitem(method_name, *args, **kwargs)
 
         if method_name == "__setitem__":
             return ArrayApi.generic_method(self, method_name, *args, **kwargs)
 
-        if method_name.startswith("__") and method_name.endswith("__"):
-            # An operator.
-            if len(args) not in (0, 1):
-                raise ValueError(
-                    f"An operator must have zero or one argument not {len(args)}."
-                )
-            if len(kwargs) not in (0, 1):
-                raise ValueError(f"Operators do not support parameters {len(kwargs)}.")
-
-            eag = eager_onnx(EagerTensor._op_impl, self.__class__, bypass_eager=True)
-            res = eag(self, *args, method_name=method_name, already_eager=True)
-            if isinstance(res, tuple) and len(res) == 1:
-                return res[0]
-            return res
-
         if method_name in {"mean", "sum", "min", "max", "prod"}:
-            # ReduceFunction
-            if len(args) not in (0, 1):
-                raise ValueError(
-                    f"An operator must have zero or one argument not {len(args)}."
-                )
-
-            if "axis" in kwargs:
-                axes = kwargs["axis"]
-                del kwargs["axis"]
-            else:
-                axes = None
-            if axes is None:
-                eag = eager_onnx(
-                    EagerTensor._reduce_impl_noaxes, self.__class__, bypass_eager=True
-                )
-                res = eag(self, method_name=method_name, already_eager=True, **kwargs)
-            else:
-                eag = eager_onnx(
-                    EagerTensor._reduce_impl, self.__class__, bypass_eager=True
-                )
-                res = eag(
-                    self, axes, method_name=method_name, already_eager=True, **kwargs
-                )
-            if isinstance(res, tuple) and len(res) == 1:
-                return res[0]
-            return res
+            return self._generic_method_reduce(method_name, *args, **kwargs)
+
+        if method_name == "astype":
+            return self._generic_method_astype(method_name, *args, **kwargs)
+
+        if method_name.startswith("__") and method_name.endswith("__"):
+            return self._generic_method_operator(method_name, *args, **kwargs)
 
         return ArrayApi.generic_method(self, method_name, *args, **kwargs)
