def check(shape, index_ndim, mode, order): multi_index = np.unravel_index(np.arange(np.product(shape)), shape, order=order) # create some invalid indices to test the mode if mode in ("wrap", "clip"): multi_index = (multi_index[0] - 1, ) + multi_index[1:] # test with scalars and higher-dimensional indices if index_ndim == 0: multi_index = tuple(i[-1] for i in multi_index) elif index_ndim == 2: multi_index = tuple(i[:, np.newaxis] for i in multi_index) multi_index_symb = [theano.shared(i) for i in multi_index] # reference result ref = np.ravel_multi_index(multi_index, shape, mode, order) def fn(mi, s): return function([], ravel_multi_index(mi, s, mode, order)) # shape given as a tuple f_array_tuple = fn(multi_index, shape) f_symb_tuple = fn(multi_index_symb, shape) np.testing.assert_equal(ref, f_array_tuple()) np.testing.assert_equal(ref, f_symb_tuple()) # shape given as an array shape_array = np.array(shape) f_array_array = fn(multi_index, shape_array) np.testing.assert_equal(ref, f_array_array()) # shape given as a theano variable shape_symb = theano.shared(shape_array) f_array_symb = fn(multi_index, shape_symb) np.testing.assert_equal(ref, f_array_symb()) # shape testing self._compile_and_check( [], [ravel_multi_index(multi_index, shape_symb, mode, order)], [], RavelMultiIndex, )
def check(shape, index_ndim, mode, order): multi_index = np.unravel_index(np.arange(np.product(shape)), shape, order=order) # create some invalid indices to test the mode if mode in ('wrap', 'clip'): multi_index = (multi_index[0] - 1,) + multi_index[1:] # test with scalars and higher-dimensional indices if index_ndim == 0: multi_index = tuple(i[-1] for i in multi_index) elif index_ndim == 2: multi_index = tuple(i[:, np.newaxis] for i in multi_index) multi_index_symb = [theano.shared(i) for i in multi_index] # reference result ref = np.ravel_multi_index(multi_index, shape, mode, order) def fn(mi, s): return function([], ravel_multi_index(mi, s, mode, order)) # shape given as a tuple f_array_tuple = fn(multi_index, shape) f_symb_tuple = fn(multi_index_symb, shape) np.testing.assert_equal(ref, f_array_tuple()) np.testing.assert_equal(ref, f_symb_tuple()) # shape given as an array shape_array = np.array(shape) f_array_array = fn(multi_index, shape_array) np.testing.assert_equal(ref, f_array_array()) # shape given as a theano variable shape_symb = theano.shared(shape_array) f_array_symb = fn(multi_index, shape_symb) np.testing.assert_equal(ref, f_array_symb()) # shape testing self._compile_and_check([], [ravel_multi_index(multi_index, shape_symb, mode, order)], [], RavelMultiIndex)
def fn(mi, s): return function([], ravel_multi_index(mi, s, mode, order))
def test_ravel_multi_index(self): def check(shape, index_ndim, mode, order): multi_index = np.unravel_index(np.arange(np.product(shape)), shape, order=order) # create some invalid indices to test the mode if mode in ("wrap", "clip"): multi_index = (multi_index[0] - 1, ) + multi_index[1:] # test with scalars and higher-dimensional indices if index_ndim == 0: multi_index = tuple(i[-1] for i in multi_index) elif index_ndim == 2: multi_index = tuple(i[:, np.newaxis] for i in multi_index) multi_index_symb = [theano.shared(i) for i in multi_index] # reference result ref = np.ravel_multi_index(multi_index, shape, mode, order) def fn(mi, s): return function([], ravel_multi_index(mi, s, mode, order)) # shape given as a tuple f_array_tuple = fn(multi_index, shape) f_symb_tuple = fn(multi_index_symb, shape) np.testing.assert_equal(ref, f_array_tuple()) np.testing.assert_equal(ref, f_symb_tuple()) # shape given as an array shape_array = np.array(shape) f_array_array = fn(multi_index, shape_array) np.testing.assert_equal(ref, f_array_array()) # shape given as a theano variable shape_symb = theano.shared(shape_array) f_array_symb = fn(multi_index, shape_symb) np.testing.assert_equal(ref, f_array_symb()) # shape testing self._compile_and_check( [], [ravel_multi_index(multi_index, shape_symb, mode, order)], [], RavelMultiIndex, ) for mode in ("raise", "wrap", "clip"): for order in ("C", "F"): for index_ndim in (0, 1, 2): check((3, ), index_ndim, mode, order) check((3, 4), index_ndim, mode, order) check((3, 4, 5), index_ndim, mode, order) # must provide integers with pytest.raises(TypeError): ravel_multi_index( (theano.tensor.fvector(), theano.tensor.ivector()), (3, 4)) with pytest.raises(TypeError): ravel_multi_index(((3, 4), theano.tensor.ivector()), (3.4, 3.2)) # dims must be a 1D sequence with pytest.raises(TypeError): ravel_multi_index(((3, 4), ), ((3, 4), ))