Skip to content

GitLab

Switch branch/tag
Find file Normal viewHistoryPermalink
test_numpy_interop.py 4.13 KB
Newer Older
root's avatar
clone code from github  
root committed
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
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150 
import unittest

import numpy
import pytest

import cupy
from cupy import testing
import cupyx

try:
    import scipy.sparse
    scipy_available = True
except ImportError:
    scipy_available = False


class TestGetArrayModule(unittest.TestCase):

    def test_get_array_module_1(self):
        n1 = numpy.array([2], numpy.float32)
        c1 = cupy.array([2], numpy.float32)
        csr1 = cupyx.scipy.sparse.csr_matrix((5, 3), dtype=numpy.float32)

        assert numpy is cupy.get_array_module()
        assert numpy is cupy.get_array_module(n1)
        assert cupy is cupy.get_array_module(c1)
        assert cupy is cupy.get_array_module(csr1)

        assert numpy is cupy.get_array_module(n1, n1)
        assert cupy is cupy.get_array_module(c1, c1)
        assert cupy is cupy.get_array_module(csr1, csr1)

        assert cupy is cupy.get_array_module(n1, csr1)
        assert cupy is cupy.get_array_module(csr1, n1)
        assert cupy is cupy.get_array_module(c1, n1)
        assert cupy is cupy.get_array_module(n1, c1)
        assert cupy is cupy.get_array_module(c1, csr1)
        assert cupy is cupy.get_array_module(csr1, c1)

        if scipy_available:
            csrn1 = scipy.sparse.csr_matrix((5, 3), dtype=numpy.float32)

            assert numpy is cupy.get_array_module(csrn1)
            assert cupy is cupy.get_array_module(csrn1, csr1)
            assert cupy is cupy.get_array_module(csr1, csrn1)
            assert cupy is cupy.get_array_module(c1, csrn1)
            assert cupy is cupy.get_array_module(csrn1, c1)
            assert numpy is cupy.get_array_module(n1, csrn1)
            assert numpy is cupy.get_array_module(csrn1, n1)


class MockArray(numpy.lib.mixins.NDArrayOperatorsMixin):
    __array_priority__ = 20  # less than cupy.ndarray.__array_priority__

    def __array_ufunc__(self, ufunc, method, *inputs, **kwargs):
        assert method == '__call__'
        name = ufunc.__name__
        return name, inputs, kwargs


class TestArrayUfunc:

    def test_add(self):
        x = cupy.array([3, 7])
        y = MockArray()
        assert x + y == ('add', (x, y), {})
        assert y + x == ('add', (y, x), {})
        y2 = y
        y2 += x
        assert y2 == ('add', (y, x), {'out': y})
        with pytest.raises(TypeError):
            x += y

    @pytest.mark.xfail(
        reason='cupy.ndarray.__array_ufunc__ does not support gufuncs yet')
    def test_matmul(self):
        x = cupy.array([3, 7])
        y = MockArray()
        assert x @ y == ('matmul', (x, y), {})
        assert y @ x == ('matmul', (y, x), {})
        y2 = y
        y2 @= x
        assert y2 == ('matmul', (y, x), {'out': y})
        with pytest.raises(TypeError):
            x @= y

    def test_lt(self):
        x = cupy.array([3, 7])
        y = MockArray()
        assert (x < y) == ('less', (x, y), {})
        assert (y < x) == ('less', (y, x), {})


class MockArray2:
    __array_ufunc__ = None

    def __add__(self, other):
        return 'add'

    def __radd__(self, other):
        return 'radd'

    def __matmul__(self, other):
        return 'matmul'

    def __rmatmul__(self, other):
        return 'rmatmul'

    def __lt__(self, other):
        return 'lt'

    def __gt__(self, other):
        return 'gt'


class TestArrayUfuncOptout:

    def test_add(self):
        x = cupy.array([3, 7])
        y = MockArray2()
        assert x + y == 'radd'
        assert y + x == 'add'

    def test_matmul(self):
        x = cupy.array([3, 7])
        y = MockArray2()
        assert x @ y == 'rmatmul'
        assert y @ x == 'matmul'

    def test_lt(self):
        x = cupy.array([3, 7])
        y = MockArray2()
        assert (x < y) == 'gt'
        assert (y < x) == 'lt'


class TestAsnumpy:

    def test_asnumpy(self):
        x = testing.shaped_random((2, 3, 4), cupy, cupy.float64)
        y = cupy.asnumpy(x)
        testing.assert_array_equal(x, y)

    def test_asnumpy_out(self):
        x = testing.shaped_random((2, 3, 4), cupy, cupy.float64)
        y = cupyx.empty_like_pinned(x)
        y = cupy.asnumpy(x, out=y)
        testing.assert_array_equal(x, y)
        assert isinstance(y.base, cupy.cuda.PinnedMemoryPointer)
        assert y.base.ptr == y.ctypes.data