fix eager mode with sklearn

Author: xadupre

_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

@@ -10,8 +10,19 @@
 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()
@@ -23,11 +34,10 @@ def test_sklearn_array_api_linear_discriminant(self):
         self.assertStartsWith("EagerNumpyTensor(array([[", repr(new_x))
         with config_context(array_api_dispatch=True):
             got = ana.predict(new_x)
-        self.assertEqualArray(expected, got)
+        self.assertEqualArray(expected, got.numpy())
 
 
 if __name__ == "__main__":
-    import logging
-
-    logging.basicConfig(level=logging.DEBUG)
+    # import logging
+    # logging.basicConfig(level=logging.DEBUG)
     unittest.main(verbosity=2)

onnx_array_api/npx/npx_functions.py

@@ -526,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_jit_eager.py

@@ -145,15 +145,15 @@ def make_key(*values, **kwargs):
                     elif isinstance(sk, (int, float)):
                         res.append(("slice", sk))
                     else:
-                        raise TypeError(f"Input {iv} cannot be such tuple: {v}.")
+                        raise TypeError(f"Input {iv} cannot have such tuple: {v}.")
                 res.append(tuple(subkey))
             else:
                 raise TypeError(
                     f"Unable to build a key, input {iv} has type {type(v)}."
                 )
         if kwargs:
             for k, v in sorted(kwargs.items()):
-                if isinstance(v, (int, float, str)):
+                if isinstance(v, (int, float, str, type)):
                     res.append(k)
                     res.append(v)
                 elif isinstance(v, tuple):
@@ -307,8 +307,8 @@ def move_input_to_kwargs(
         :return: new values, new arguments
         """
         if self.input_to_kwargs_ is None:
-            if self.bypass_eager or self.f.__annotations__:
-                return values, kwargs
+            # if self.bypass_eager or self.f.__annotations__:
+            #    return values, kwargs
             raise RuntimeError(
                 f"self.input_to_kwargs_ is not initialized for function {self.f} "
                 f"from module {self.f.__module__!r}."
@@ -321,8 +321,8 @@ def move_input_to_kwargs(
             if i in self.input_to_kwargs_:
                 new_kwargs[self.input_to_kwargs_[i]] = v
             else:
-                new_values.append(values)
-        return new_values, new_kwargs
+                new_values.append(v)
+        return tuple(new_values), new_kwargs
 
     def jit_call(self, *values, **kwargs):
         """
@@ -334,15 +334,33 @@ def jit_call(self, *values, **kwargs):
         and returns the result or the results in a tuple if there are several.
         """
         self.info("+", "jit_call")
+        if self.input_to_kwargs_ is None:
+            # No jitting was ever called.
+            onx, fct = self.to_jit(*values, **kwargs)
+            if self.input_to_kwargs_ is None:
+                raise RuntimeError(
+                    f"Attribute 'input_to_kwargs_' should be set for "
+                    f"function {self.f} form module {self.f.__module__!r}."
+                )
+        else:
+            onx, fct = None, None
+
         values, kwargs = self.move_input_to_kwargs(values, kwargs)
         key = self.make_key(*values, **kwargs)
         if self.method_name_ is None and "method_name" in key:
             pos = list(key).index("method_name")
             self.method_name_ = key[pos + 1]
 
-        if key in self.versions:
+        if onx is not None:
+            # First jitting.
+            self.versions[key] = fct
+            self.onxs[key] = onx
+        elif key in self.versions:
+            # Already jitted.
             fct = self.versions[key]
         else:
+            # One version was already jitted but types or parameter
+            # are different.
             onx, fct = self.to_jit(*values, **kwargs)
             self.versions[key] = fct
             self.onxs[key] = onx
@@ -482,6 +500,9 @@ def _preprocess_constants(self, *args):
             #            f"function {self.f} from module {self.f.__module__!r}."
             #        )
             #    new_args.append(n)
+            elif isinstance(n, np.ndarray):
+                new_args.append(self.tensor_class(n))
+                modified = True
             elif isinstance(n, (int, float)):
                 new_args.append(self.tensor_class(np.array(n)))
                 modified = True
@@ -518,7 +539,8 @@ def __call__(self, *args, already_eager=False, **kwargs):
                 map(
                     lambda t: t is not None
                     and not isinstance(
-                        t, (EagerTensor, Cst, int, float, tuple, slice, type)
+                        t,
+                        (EagerTensor, Cst, int, float, tuple, slice, type, np.ndarray),
                     ),
                     args,
                 )

onnx_array_api/npx/npx_tensors.py

@@ -30,6 +30,17 @@ def const_cast(self, to: Any = None) -> "EagerTensor":
             f"{self.__class__.__name__!r}."
         )
 
+    def __iter__(self):
+        """
+        This is not implementation in the generic case.
+        This method raises an exception with a better error message.
+        """
+        raise RuntimeError(
+            "Iterators are not implemented in the generic case. "
+            "It may be enabled for the eager mode but it might fail "
+            "when a whole function is converted into ONNX."
+        )
+
     @staticmethod
     def _op_impl(*inputs, method_name=None):
         # avoids circular imports.

onnx_array_api/npx/npx_var.py

@@ -561,6 +561,16 @@ def to_onnx(
 
     # Operators
 
+    def __iter__(self):
+        """
+        This is not implementation in the generic case.
+        This method raises an exception with a better error message.
+        """
+        raise RuntimeError(
+            "Iterators are not implemented in the generic case. "
+            "Every function using them cannot be converted into ONNX."
+        )
+
     def _binary_op(self, ov: "Var", op_name: str, **kwargs) -> "Var":
         var = Var.get_cst_var()[1]
         if isinstance(ov, (int, float, np.ndarray, Cst)):