onnx-array-api/_unittests/ut_npx/test_sklearn_array_api.py at 0471c35b65d09ed0fbb014513621ac3f698630f4 · sdpython/onnx-array-api · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
import unittest
import numpy as np
from packaging.version import Version
from onnx.defs import onnx_opset_version
from sklearn import config_context, __version__ as sklearn_version
from sklearn.discriminant_analysis import LinearDiscriminantAnalysis
from onnx_array_api.ext_test_case import ExtTestCase, ignore_warnings
from onnx_array_api.npx.npx_numpy_tensors import EagerNumpyTensor

DEFAULT_OPSET = onnx_opset_version()


class TestSklearnArrayAPI(ExtTestCase):
    @unittest.skipIf(
        Version(sklearn_version) <= Version("1.2.2"),
        reason="reshape ArrayAPI not followed",
    )
    @ignore_warnings(DeprecationWarning)
    @unittest.skip("not maintained")
    def test_sklearn_array_api_linear_discriminant(self):
        X = np.array(
            [[-1, -1], [-2, -1], [-3, -2], [1, 1], [2, 1], [3, 2]], dtype=np.float64
        )
        y = np.array([1, 1, 1, 2, 2, 2], dtype=np.int64)
        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):
            # 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())

    @unittest.skipIf(
        Version(sklearn_version) <= Version("1.2.2"),
        reason="reshape ArrayAPI not followed",
    )
    @ignore_warnings(DeprecationWarning)
    @unittest.skip("not maintained")
    def test_sklearn_array_api_linear_discriminant_float32(self):
        X = np.array(
            [[-1, -1], [-2, -1], [-3, -2], [1, 1], [2, 1], [3, 2]], dtype=np.float32
        )
        y = np.array([1, 1, 1, 2, 2, 2], dtype=np.int64)
        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):
            # 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())


if __name__ == "__main__":
    # import logging

    # logging.basicConfig(level=logging.DEBUG)
    unittest.main(verbosity=2)