new udpates

sdpython · sdpython · Jun 10, 2023 · Jun 5, 2023 · Jun 5, 2023 · Jun 5, 2023
commit 550d0dc6fe061d7075293afd7aae18542c2f4f40
diff --git a/_unittests/ut_npx/test_sklearn_array_api.py b/_unittests/ut_npx/test_sklearn_array_api.py
@@ -27,6 +27,8 @@ def test_sklearn_array_api_linear_discriminant(self):
         new_x = EagerNumpyTensor(X)
         self.assertStartsWith("EagerNumpyTensor(array([[", repr(new_x))
         with config_context(array_api_dispatch=True):
+            # It fails if scikit-learn <= 1.2.2 because the ArrayAPI
+            # is not strictly applied.
             got = ana.predict(new_x)
         self.assertEqualArray(expected, got.numpy())
 

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(
     else:
         raise RuntimeError(f"Unexpected type {type(a)} for the first input.")
     if dtype is not None:
-        vt = v.astype(dtype=dtype)
+        vt = v.astype(dtype)
     else:
         vt = v
     return vt
diff --git a/onnx_array_api/npx/npx_jit_eager.py b/onnx_array_api/npx/npx_jit_eager.py
@@ -131,7 +131,7 @@ def make_key(*values, **kwargs):
         for iv, v in enumerate(values):
             if isinstance(v, (Var, EagerTensor, JitTensor)):
                 res.append(v.key)
-            elif isinstance(v, (int, float)):
+            elif isinstance(v, (int, float, DType)):
                 res.append(v)
             elif isinstance(v, slice):
                 res.append(("slice", v.start, v.stop, v.step))
@@ -344,7 +344,11 @@ def jit_call(self, *values, **kwargs):
         self.info("+", "jit_call")
         if self.input_to_kwargs_ is None:
             # No jitting was ever called.
-            onx, fct = self.to_jit(*values, **kwargs)
+            try:
+                onx, fct = self.to_jit(*values, **kwargs)
+            except Exception as e:
+                raise RuntimeError(f"ERROR with self.f={self.f}, "
+                            f"values={values!r}, kwargs={kwargs!r}") from e
             if self.input_to_kwargs_ is None:
                 raise RuntimeError(
                     f"Attribute 'input_to_kwargs_' should be set for "

diff --git a/onnx_array_api/npx/npx_tensors.py b/onnx_array_api/npx/npx_tensors.py
@@ -1,9 +1,9 @@
-from typing import Any
+from typing import Any, Union
 
 import numpy as np
 from onnx.helper import np_dtype_to_tensor_dtype
 
-from .npx_types import DType, OptParType
+from .npx_types import DType, ParType
 from .npx_array_api import BaseArrayApi, ArrayApiError
 
 
@@ -74,7 +74,7 @@ def _getitem_impl_var(obj, index, method_name=None):
         return meth(obj, index)
 
     @staticmethod
-    def _astype_impl(x, dtype: OptParType[DType] = None, method_name=None):
+    def _astype_impl(x, dtype: ParType[DType], method_name=None):
         # avoids circular imports.
         if dtype is None:
             raise ValueError("dtype cannot be None.")
@@ -182,18 +182,15 @@ def _np_dtype_to_tensor_dtype(dtype):
             dtype = np.dtype("float64")
         return np_dtype_to_tensor_dtype(dtype)
 
-    def _generic_method_astype(self, method_name, *args: Any, **kwargs: Any) -> Any:
+    def _generic_method_astype(self, method_name, dtype: Union[DType, "Var"], **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], (DType, Var))
-            else self._np_dtype_to_tensor_dtype(args[0])
+            dtype
+            if isinstance(dtype, (DType, Var))
+            else self._np_dtype_to_tensor_dtype(dtype)
         )
         eag = eager_onnx(EagerTensor._astype_impl, self.__class__, bypass_eager=True)
         res = eag(self, dtype, method_name=method_name, already_eager=True, **kwargs)