1# Licensed to the Apache Software Foundation (ASF) under one
2# or more contributor license agreements.  See the NOTICE file
3# distributed with this work for additional information
4# regarding copyright ownership.  The ASF licenses this file
5# to you under the Apache License, Version 2.0 (the
6# "License"); you may not use this file except in compliance
7# with the License.  You may obtain a copy of the License at
8#
9#   http://www.apache.org/licenses/LICENSE-2.0
10#
11# Unless required by applicable law or agreed to in writing,
12# software distributed under the License is distributed on an
13# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
14# KIND, either express or implied.  See the License for the
15# specific language governing permissions and limitations
16# under the License.
17
18import os
19import mxnet as mx
20import numpy as np
21import pickle as pkl
22
23
24def _np_reduce(dat, axis, keepdims, numpy_reduce_func):
25    if isinstance(axis, int):
26        axis = [axis]
27    else:
28        axis = list(axis) if axis is not None else range(len(dat.shape))
29    ret = dat
30    for i in reversed(sorted(axis)):
31        ret = numpy_reduce_func(ret, axis=i)
32    if keepdims:
33        keepdims_shape = list(dat.shape)
34        for i in axis:
35            keepdims_shape[i] = 1
36        ret = ret.reshape(tuple(keepdims_shape))
37    return ret
38
39
40def reldiff(a, b):
41    diff = np.abs(a - b)
42    norm = np.abs(a)
43    reldiff = np.max(diff  / (norm + 1e-7))
44    return reldiff
45
46
47def same(a, b):
48    return np.sum(a != b) == 0
49
50
51def check_with_uniform(uf, arg_shapes, dim=None, npuf=None, rmin=-10, type_list=[np.float32]):
52    """check function consistency with uniform random numbers"""
53    if isinstance(arg_shapes, int):
54        assert dim
55        shape = tuple(np.random.randint(1, int(1000**(1.0/dim)), size=dim))
56        arg_shapes = [shape] * arg_shapes
57    for dtype in type_list:
58        ndarray_arg = []
59        numpy_arg = []
60        for s in arg_shapes:
61            npy = np.random.uniform(rmin, 10, s).astype(dtype)
62            narr = mx.nd.array(npy, dtype=dtype)
63            ndarray_arg.append(narr)
64            numpy_arg.append(npy)
65        out1 = uf(*ndarray_arg)
66        if npuf is None:
67            out2 = uf(*numpy_arg).astype(dtype)
68        else:
69            out2 = npuf(*numpy_arg).astype(dtype)
70
71        assert out1.shape == out2.shape
72        if isinstance(out1, mx.nd.NDArray):
73            out1 = out1.asnumpy()
74        if dtype == np.float16:
75            assert reldiff(out1, out2) < 2e-3
76        else:
77            assert reldiff(out1, out2) < 1e-6
78
79
80def random_ndarray(dim):
81    shape = tuple(np.random.randint(1, int(1000**(1.0/dim)), size=dim))
82    data = mx.nd.array(np.random.uniform(-10, 10, shape))
83    return data
84
85
86def test_ndarray_elementwise():
87    np.random.seed(0)
88    nrepeat = 10
89    maxdim = 4
90    all_type = [np.float32, np.float64, np.float16, np.uint8, np.int32]
91    real_type = [np.float32, np.float64, np.float16]
92    for repeat in range(nrepeat):
93        for dim in range(1, maxdim):
94            check_with_uniform(lambda x, y: x + y, 2, dim, type_list=all_type)
95            check_with_uniform(lambda x, y: x - y, 2, dim, type_list=all_type)
96            check_with_uniform(lambda x, y: x * y, 2, dim, type_list=all_type)
97            check_with_uniform(lambda x, y: x / y, 2, dim, type_list=real_type)
98            check_with_uniform(lambda x, y: x / y, 2, dim, rmin=1, type_list=all_type)
99            check_with_uniform(mx.nd.sqrt, 1, dim, np.sqrt, rmin=0)
100            check_with_uniform(mx.nd.square, 1, dim, np.square, rmin=0)
101            check_with_uniform(lambda x: mx.nd.norm(x).asscalar(), 1, dim, np.linalg.norm)
102
103
104def test_ndarray_negate():
105    npy = np.random.uniform(-10, 10, (2,3,4))
106    arr = mx.nd.array(npy)
107    assert reldiff(npy, arr.asnumpy()) < 1e-6
108    assert reldiff(-npy, (-arr).asnumpy()) < 1e-6
109
110    # a final check to make sure the negation (-) is not implemented
111    # as inplace operation, so the contents of arr does not change after
112    # we compute (-arr)
113    assert reldiff(npy, arr.asnumpy()) < 1e-6
114
115
116def test_ndarray_choose():
117    shape = (100, 20)
118    npy = np.arange(np.prod(shape)).reshape(shape)
119    arr = mx.nd.array(npy)
120    nrepeat = 3
121    for repeat in range(nrepeat):
122        indices = np.random.randint(shape[1], size=shape[0])
123        assert same(npy[np.arange(shape[0]), indices],
124                    mx.nd.choose_element_0index(arr, mx.nd.array(indices)).asnumpy())
125
126
127def test_ndarray_fill():
128    shape = (100, 20)
129    npy = np.arange(np.prod(shape)).reshape(shape)
130    arr = mx.nd.array(npy)
131    new_npy = npy.copy()
132    nrepeat = 3
133    for repeat in range(nrepeat):
134        indices = np.random.randint(shape[1], size=shape[0])
135        val = np.random.randint(shape[1], size=shape[0])
136        new_npy[:] = npy
137        new_npy[np.arange(shape[0]), indices] = val
138        assert same(new_npy,
139                    mx.nd.fill_element_0index(arr, mx.nd.array(val), mx.nd.array(indices)).asnumpy())
140
141
142def test_ndarray_onehot():
143    shape = (100, 20)
144    npy = np.arange(np.prod(shape)).reshape(shape)
145    arr = mx.nd.array(npy)
146    nrepeat = 3
147    for repeat in range(nrepeat):
148        indices = np.random.randint(shape[1], size=shape[0])
149        npy[:] = 0.0
150        npy[np.arange(shape[0]), indices] = 1.0
151        mx.nd.onehot_encode(mx.nd.array(indices), out=arr)
152        assert same(npy, arr.asnumpy())
153
154
155def test_ndarray_copy():
156    c = mx.nd.array(np.random.uniform(-10, 10, (10, 10)))
157    d = c.copyto(mx.Context('cpu', 0))
158    assert np.sum(np.abs(c.asnumpy() != d.asnumpy())) == 0.0
159
160
161def test_ndarray_scalar():
162    c = mx.nd.empty((10,10))
163    d = mx.nd.empty((10,10))
164    c[:] = 0.5
165    d[:] = 1.0
166    d -= c * 2 / 3 * 6.0
167    c += 0.5
168    assert(np.sum(c.asnumpy()) - 100 < 1e-5)
169    assert(np.sum(d.asnumpy()) + 100 < 1e-5)
170    c[:] = 2
171    assert(np.sum(c.asnumpy()) - 200 < 1e-5)
172    d = -c + 2
173    assert(np.sum(d.asnumpy()) < 1e-5)
174
175
176def test_ndarray_pickle():
177    np.random.seed(0)
178    maxdim = 5
179    nrepeat = 10
180    for repeat in range(nrepeat):
181        for dim in range(1, maxdim):
182            a = random_ndarray(dim)
183            b = mx.nd.empty(a.shape)
184            a[:] = np.random.uniform(-10, 10, a.shape)
185            b[:] = np.random.uniform(-10, 10, a.shape)
186            a = a + b
187            data = pkl.dumps(a)
188            a2 = pkl.loads(data)
189            assert np.sum(a.asnumpy() != a2.asnumpy()) == 0
190
191
192def test_ndarray_saveload():
193    np.random.seed(0)
194    maxdim = 5
195    nrepeat = 10
196    fname = 'tmp_list.bin'
197    for repeat in range(nrepeat):
198        data = []
199        for i in range(10):
200            data.append(random_ndarray(np.random.randint(1, 5)))
201        mx.nd.save(fname, data)
202        data2 = mx.nd.load(fname)
203        assert len(data) == len(data2)
204        for x, y in zip(data, data2):
205            assert np.sum(x.asnumpy() != y.asnumpy()) == 0
206        dmap = {'ndarray xx %s' % i : x for i, x in enumerate(data)}
207        mx.nd.save(fname, dmap)
208        dmap2 = mx.nd.load(fname)
209        assert len(dmap2) == len(dmap)
210        for k, x in dmap.items():
211            y = dmap2[k]
212            assert np.sum(x.asnumpy() != y.asnumpy()) == 0
213    os.remove(fname)
214
215
216def test_ndarray_slice():
217    shape = (10,)
218    A = mx.nd.array(np.random.uniform(-10, 10, shape))
219    A2 = A.asnumpy()
220    assert same(A[3:8].asnumpy(), A2[3:8])
221    A2[3:8] *= 10;
222    A[3:8] = A2[3:8]
223    assert same(A[3:8].asnumpy(), A2[3:8])
224
225
226def test_ndarray_slice_along_axis():
227    arr = mx.nd.array(np.random.uniform(-10, 10, (3, 4, 2, 3)))
228    sub_arr = arr.slice(begin=(None, 1), end=(None, 3))
229
230    # test we sliced correctly
231    assert same(arr.asnumpy()[:, 1:3, :, :], sub_arr.asnumpy())
232
233    # test that slice is copy, instead of shared memory
234    sub_arr[:] = 0
235    assert not same(arr.asnumpy()[:, 1:3, :, :], sub_arr.asnumpy())
236
237
238def test_clip():
239    shape = (10,)
240    A = mx.random.uniform(-10, 10, shape)
241    B = mx.nd.clip(A, -2, 2)
242    B1 = B.asnumpy()
243    for i in range(shape[0]):
244        assert B1[i] >= -2
245        assert B1[i] <= 2
246
247
248def test_dot():
249    a = np.random.uniform(-3, 3, (3, 4))
250    b = np.random.uniform(-3, 3, (4, 5))
251    c = np.dot(a, b)
252    A = mx.nd.array(a)
253    B = mx.nd.array(b)
254    C = mx.nd.dot(A, B)
255    assert reldiff(c, C.asnumpy()) < 1e-5
256
257
258def test_reduce():
259    sample_num = 200
260
261    def test_reduce_inner(numpy_reduce_func, nd_reduce_func):
262        for i in range(sample_num):
263            ndim = np.random.randint(1, 6)
264            shape = np.random.randint(1, 11, size=ndim)
265            axis_flags = np.random.randint(0, 2, size=ndim)
266            axes = []
267            for (axis, flag) in enumerate(axis_flags):
268                if flag:
269                    axes.append(axis)
270            keepdims = np.random.randint(0, 2)
271            dat = np.random.rand(*shape) - 0.5
272            if 0 == len(axes):
273                axes = tuple(range(ndim))
274            else:
275                axes = tuple(axes)
276            numpy_ret = numpy_reduce_func(dat, axis=axes, keepdims=keepdims)
277
278            ndarray_ret = nd_reduce_func(mx.nd.array(dat), axis=axes, keepdims=keepdims)
279            if type(ndarray_ret) is mx.ndarray.NDArray:
280                ndarray_ret = ndarray_ret.asnumpy()
281            assert (ndarray_ret.shape == numpy_ret.shape) or \
282                   (ndarray_ret.shape == (1,) and numpy_ret.shape == ()), "nd:%s, numpy:%s" \
283                                                         %(ndarray_ret.shape, numpy_ret.shape)
284            err = np.square(ndarray_ret - numpy_ret).mean()
285            assert err < 1E-4
286    test_reduce_inner(lambda data, axis, keepdims:_np_reduce(data, axis, keepdims, np.sum),
287                      mx.nd.sum)
288    test_reduce_inner(lambda data, axis, keepdims:_np_reduce(data, axis, keepdims, np.max),
289                      mx.nd.max)
290    test_reduce_inner(lambda data, axis, keepdims:_np_reduce(data, axis, keepdims, np.min),
291                      mx.nd.min)
292
293
294def test_broadcast():
295    sample_num = 1000
296
297    def test_broadcast_to():
298        for i in range(sample_num):
299            ndim = np.random.randint(1, 6)
300            target_shape = np.random.randint(1, 11, size=ndim)
301            shape = target_shape.copy()
302            axis_flags = np.random.randint(0, 2, size=ndim)
303            axes = []
304            for (axis, flag) in enumerate(axis_flags):
305                if flag:
306                    shape[axis] = 1
307            dat = np.random.rand(*shape) - 0.5
308            numpy_ret = dat
309            ndarray_ret = mx.nd.array(dat).broadcast_to(shape=target_shape)
310            if type(ndarray_ret) is mx.ndarray.NDArray:
311                ndarray_ret = ndarray_ret.asnumpy()
312            assert (ndarray_ret.shape == target_shape).all()
313            err = np.square(ndarray_ret - numpy_ret).mean()
314            assert err < 1E-8
315    test_broadcast_to()
316
317
318if __name__ == '__main__':
319    mx.profiler.set_config(profile_all=True, filename='profile_ndarray.json')
320    mx.profiler.set_state('run')
321    test_ndarray_slice_along_axis()
322    test_broadcast()
323    test_ndarray_elementwise()
324    test_ndarray_slice()
325    test_ndarray_pickle()
326    test_ndarray_saveload()
327    test_ndarray_copy()
328    test_ndarray_negate()
329    test_ndarray_scalar()
330    test_clip()
331    test_dot()
332    test_ndarray_choose()
333    test_ndarray_onehot()
334    test_ndarray_fill()
335    test_reduce()
336    mx.profiler.set_state('stop')
337