Check Array API with onnxruntime

sdpython · sdpython · Jun 5, 2023 · Mar 20, 2023 · Jun 3, 2023 · Jun 3, 2023
commit 602fd9eb9de973b26e5a0744a5b0bfaaf5ce346f
diff --git a/_unittests/ut_ort/test_sklearn_array_api_ort.py b/_unittests/ut_ort/test_sklearn_array_api_ort.py
@@ -0,0 +1,46 @@
+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.ort.ort_tensors import EagerOrtTensor, OrtTensor
+
+
+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 TestSklearnArrayAPIOrt(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 = EagerOrtTensor(OrtTensor.from_array(X))
+        self.assertEqual(new_x.device_name, "Cpu")
+        self.assertStartsWith(
+            "EagerOrtTensor(OrtTensor.from_array(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)
diff --git a/onnx_array_api/npx/npx_jit_eager.py b/onnx_array_api/npx/npx_jit_eager.py
@@ -370,8 +370,11 @@ def jit_call(self, *values, **kwargs):
             raise RuntimeError(
                 f"Unable to run function for key={key!r}, "
                 f"types={[type(x) for x in values]}, "
+                f"dtypes={[x.dtype for x in values]}, "
+                f"shapes={[x.shape for x in values]}, "
                 f"kwargs={kwargs}, "
                 f"self.input_to_kwargs_={self.input_to_kwargs_}, "
+                f"f={self.f} from module {self.f.__module__!r} "
                 f"onnx={self.onxs[key]}."
             ) from e
         self.info("-", "jit_call")

diff --git a/onnx_array_api/npx/npx_numpy_tensors.py b/onnx_array_api/npx/npx_numpy_tensors.py
@@ -162,12 +162,6 @@ def get_ir_version(cls, ir_version):
         """
         return ir_version
 
-    def const_cast(self, to: Any = None) -> "EagerTensor":
-        """
-        Casts a constant without any ONNX conversion.
-        """
-        return self.__class__(self._tensor.astype(to))
-
     # The class should support whatever Var supports.
     # This part is not yet complete.
 

diff --git a/onnx_array_api/npx/npx_tensors.py b/onnx_array_api/npx/npx_tensors.py
@@ -21,15 +21,6 @@ class EagerTensor(ArrayApi):
     :class:`ArrayApi`.
     """
 
-    def const_cast(self, to: Any = None) -> "EagerTensor":
-        """
-        Casts a constant without any ONNX conversion.
-        """
-        raise NotImplementedError(
-            f"Method 'const_cast' must be overwritten in class "
-            f"{self.__class__.__name__!r}."
-        )
-
     def __iter__(self):
         """
         This is not implementation in the generic case.
@@ -141,7 +132,7 @@ def _generic_method_operator(self, method_name, *args: Any, **kwargs: Any) -> An
         new_args = []
         for a in args:
             if isinstance(a, np.ndarray):
-                new_args.append(self.__class__(a).const_cast(self.dtype))
+                new_args.append(self.__class__(a.astype(self.dtype)))
             else:
                 new_args.append(a)
 

diff --git a/onnx_array_api/ort/ort_tensors.py b/onnx_array_api/ort/ort_tensors.py
@@ -3,6 +3,7 @@
 import numpy as np
 from onnx import ModelProto, TensorProto
 from onnx.defs import onnx_opset_version
+from onnx.helper import tensor_dtype_to_np_dtype
 from onnxruntime import InferenceSession, RunOptions, get_available_providers
 from onnxruntime.capi._pybind_state import OrtDevice as C_OrtDevice
 from onnxruntime.capi._pybind_state import OrtMemType
@@ -128,9 +129,23 @@ def __init__(self, tensor: Union[C_OrtValue, "OrtTensor"]):
             self._tensor = tensor
         elif isinstance(tensor, OrtTensor):
             self._tensor = tensor._tensor
+        elif isinstance(tensor, np.ndarray):
+            self._tensor = C_OrtValue.ortvalue_from_numpy(tensor, OrtTensor.CPU)
         else:
             raise ValueError(f"An OrtValue is expected not {type(tensor)}.")
 
+    def __repr__(self) -> str:
+        "usual"
+        return f"{self.__class__.__name__}(OrtTensor.from_array({self.numpy()!r}))"
+
+    @property
+    def device_name(self):
+        return self._tensor.device_name()
+
+    @property
+    def ndim(self):
+        return len(self.shape)
+
     @property
     def shape(self) -> Tuple[int, ...]:
         "Returns the shape of the tensor."
@@ -139,7 +154,7 @@ def shape(self) -> Tuple[int, ...]:
     @property
     def dtype(self) -> Any:
         "Returns the element type of this tensor."
-        return self._tensor.element_type()
+        return tensor_dtype_to_np_dtype(self._tensor.element_type())
 
     @property
     def key(self) -> Any: