CoolFace
Apppublic

Aluode/PerceptionLabPortable

sourceHugging Faceupdated 9mo agoView on Hugging Face
0likes
test_cython_blas.py251 linesDownload Raw Back to tests
1import numpy as np
2import pytest
3
4from sklearn.utils._cython_blas import (
5    BLAS_Order,
6    BLAS_Trans,
7    _asum_memview,
8    _axpy_memview,
9    _copy_memview,
10    _dot_memview,
11    _gemm_memview,
12    _gemv_memview,
13    _ger_memview,
14    _nrm2_memview,
15    _rot_memview,
16    _rotg_memview,
17    _scal_memview,
18)
19from sklearn.utils._testing import assert_allclose
20
21
22def _numpy_to_cython(dtype):
23    cython = pytest.importorskip("cython")
24    if dtype == np.float32:
25        return cython.float
26    elif dtype == np.float64:
27        return cython.double
28
29
30RTOL = {np.float32: 1e-6, np.float64: 1e-12}
31ORDER = {BLAS_Order.RowMajor: "C", BLAS_Order.ColMajor: "F"}
32
33
34def _no_op(x):
35    return x
36
37
38@pytest.mark.parametrize("dtype", [np.float32, np.float64])
39def test_dot(dtype):
40    dot = _dot_memview[_numpy_to_cython(dtype)]
41
42    rng = np.random.RandomState(0)
43    x = rng.random_sample(10).astype(dtype, copy=False)
44    y = rng.random_sample(10).astype(dtype, copy=False)
45
46    expected = x.dot(y)
47    actual = dot(x, y)
48
49    assert_allclose(actual, expected, rtol=RTOL[dtype])
50
51
52@pytest.mark.parametrize("dtype", [np.float32, np.float64])
53def test_asum(dtype):
54    asum = _asum_memview[_numpy_to_cython(dtype)]
55
56    rng = np.random.RandomState(0)
57    x = rng.random_sample(10).astype(dtype, copy=False)
58
59    expected = np.abs(x).sum()
60    actual = asum(x)
61
62    assert_allclose(actual, expected, rtol=RTOL[dtype])
63
64
65@pytest.mark.parametrize("dtype", [np.float32, np.float64])
66def test_axpy(dtype):
67    axpy = _axpy_memview[_numpy_to_cython(dtype)]
68
69    rng = np.random.RandomState(0)
70    x = rng.random_sample(10).astype(dtype, copy=False)
71    y = rng.random_sample(10).astype(dtype, copy=False)
72    alpha = 2.5
73
74    expected = alpha * x + y
75    axpy(alpha, x, y)
76
77    assert_allclose(y, expected, rtol=RTOL[dtype])
78
79
80@pytest.mark.parametrize("dtype", [np.float32, np.float64])
81def test_nrm2(dtype):
82    nrm2 = _nrm2_memview[_numpy_to_cython(dtype)]
83
84    rng = np.random.RandomState(0)
85    x = rng.random_sample(10).astype(dtype, copy=False)
86
87    expected = np.linalg.norm(x)
88    actual = nrm2(x)
89
90    assert_allclose(actual, expected, rtol=RTOL[dtype])
91
92
93@pytest.mark.parametrize("dtype", [np.float32, np.float64])
94def test_copy(dtype):
95    copy = _copy_memview[_numpy_to_cython(dtype)]
96
97    rng = np.random.RandomState(0)
98    x = rng.random_sample(10).astype(dtype, copy=False)
99    y = np.empty_like(x)
100
101    expected = x.copy()
102    copy(x, y)
103
104    assert_allclose(y, expected, rtol=RTOL[dtype])
105
106
107@pytest.mark.parametrize("dtype", [np.float32, np.float64])
108def test_scal(dtype):
109    scal = _scal_memview[_numpy_to_cython(dtype)]
110
111    rng = np.random.RandomState(0)
112    x = rng.random_sample(10).astype(dtype, copy=False)
113    alpha = 2.5
114
115    expected = alpha * x
116    scal(alpha, x)
117
118    assert_allclose(x, expected, rtol=RTOL[dtype])
119
120
121@pytest.mark.parametrize("dtype", [np.float32, np.float64])
122def test_rotg(dtype):
123    rotg = _rotg_memview[_numpy_to_cython(dtype)]
124
125    rng = np.random.RandomState(0)
126    a = dtype(rng.randn())
127    b = dtype(rng.randn())
128    c, s = 0.0, 0.0
129
130    def expected_rotg(a, b):
131        roe = a if abs(a) > abs(b) else b
132        if a == 0 and b == 0:
133            c, s, r, z = (1, 0, 0, 0)
134        else:
135            r = np.sqrt(a**2 + b**2) * (1 if roe >= 0 else -1)
136            c, s = a / r, b / r
137            z = s if roe == a else (1 if c == 0 else 1 / c)
138        return r, z, c, s
139
140    expected = expected_rotg(a, b)
141    actual = rotg(a, b, c, s)
142
143    assert_allclose(actual, expected, rtol=RTOL[dtype])
144
145
146@pytest.mark.parametrize("dtype", [np.float32, np.float64])
147def test_rot(dtype):
148    rot = _rot_memview[_numpy_to_cython(dtype)]
149
150    rng = np.random.RandomState(0)
151    x = rng.random_sample(10).astype(dtype, copy=False)
152    y = rng.random_sample(10).astype(dtype, copy=False)
153    c = dtype(rng.randn())
154    s = dtype(rng.randn())
155
156    expected_x = c * x + s * y
157    expected_y = c * y - s * x
158
159    rot(x, y, c, s)
160
161    assert_allclose(x, expected_x)
162    assert_allclose(y, expected_y)
163
164
165@pytest.mark.parametrize("dtype", [np.float32, np.float64])
166@pytest.mark.parametrize(
167    "opA, transA",
168    [(_no_op, BLAS_Trans.NoTrans), (np.transpose, BLAS_Trans.Trans)],
169    ids=["NoTrans", "Trans"],
170)
171@pytest.mark.parametrize(
172    "order",
173    [BLAS_Order.RowMajor, BLAS_Order.ColMajor],
174    ids=["RowMajor", "ColMajor"],
175)
176def test_gemv(dtype, opA, transA, order):
177    gemv = _gemv_memview[_numpy_to_cython(dtype)]
178
179    rng = np.random.RandomState(0)
180    A = np.asarray(
181        opA(rng.random_sample((20, 10)).astype(dtype, copy=False)), order=ORDER[order]
182    )
183    x = rng.random_sample(10).astype(dtype, copy=False)
184    y = rng.random_sample(20).astype(dtype, copy=False)
185    alpha, beta = 2.5, -0.5
186
187    expected = alpha * opA(A).dot(x) + beta * y
188    gemv(transA, alpha, A, x, beta, y)
189
190    assert_allclose(y, expected, rtol=RTOL[dtype])
191
192
193@pytest.mark.parametrize("dtype", [np.float32, np.float64])
194@pytest.mark.parametrize(
195    "order",
196    [BLAS_Order.RowMajor, BLAS_Order.ColMajor],
197    ids=["BLAS_Order.RowMajor", "BLAS_Order.ColMajor"],
198)
199def test_ger(dtype, order):
200    ger = _ger_memview[_numpy_to_cython(dtype)]
201
202    rng = np.random.RandomState(0)
203    x = rng.random_sample(10).astype(dtype, copy=False)
204    y = rng.random_sample(20).astype(dtype, copy=False)
205    A = np.asarray(
206        rng.random_sample((10, 20)).astype(dtype, copy=False), order=ORDER[order]
207    )
208    alpha = 2.5
209
210    expected = alpha * np.outer(x, y) + A
211    ger(alpha, x, y, A)
212
213    assert_allclose(A, expected, rtol=RTOL[dtype])
214
215
216@pytest.mark.parametrize("dtype", [np.float32, np.float64])
217@pytest.mark.parametrize(
218    "opB, transB",
219    [(_no_op, BLAS_Trans.NoTrans), (np.transpose, BLAS_Trans.Trans)],
220    ids=["NoTrans", "Trans"],
221)
222@pytest.mark.parametrize(
223    "opA, transA",
224    [(_no_op, BLAS_Trans.NoTrans), (np.transpose, BLAS_Trans.Trans)],
225    ids=["NoTrans", "Trans"],
226)
227@pytest.mark.parametrize(
228    "order",
229    [BLAS_Order.RowMajor, BLAS_Order.ColMajor],
230    ids=["BLAS_Order.RowMajor", "BLAS_Order.ColMajor"],
231)
232def test_gemm(dtype, opA, transA, opB, transB, order):
233    gemm = _gemm_memview[_numpy_to_cython(dtype)]
234
235    rng = np.random.RandomState(0)
236    A = np.asarray(
237        opA(rng.random_sample((30, 10)).astype(dtype, copy=False)), order=ORDER[order]
238    )
239    B = np.asarray(
240        opB(rng.random_sample((10, 20)).astype(dtype, copy=False)), order=ORDER[order]
241    )
242    C = np.asarray(
243        rng.random_sample((30, 20)).astype(dtype, copy=False), order=ORDER[order]
244    )
245    alpha, beta = 2.5, -0.5
246
247    expected = alpha * opA(A).dot(opB(B)) + beta * C
248    gemm(transA, transB, alpha, A, B, beta, C)
249
250    assert_allclose(C, expected, rtol=RTOL[dtype])
251 
Aluode/PerceptionLabPortable · CoolFace