def check(shape, index_ndim, order):
indices = np.arange(np.product(shape))
# test with scalars and higher-dimensional indices
if index_ndim == 0:
indices = indices[-1]
elif index_ndim == 2:
indices = indices[:, np.newaxis]
indices_symb = theano.shared(indices)
# reference result
ref = np.unravel_index(indices, shape)
def fn(i, d, nd=None):
if nd is None:
return function([], unravel_index(i, d, order=order))
else:
return function([],
unravel_index(i, d, order=order, ndim=nd))
# shape given as a tuple
f_array_tuple = fn(indices, shape)
f_symb_tuple = fn(indices_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(indices, 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(indices, shape_symb, len(shape))
np.testing.assert_equal(ref, f_array_symb())
# shape given as a Shape op (unravel_index will use get_vector_length
# to infer the number of dimensions)
indexed_array = theano.shared(np.random.uniform(size=shape_array))
f_array_shape = fn(indices, indexed_array.shape)
np.testing.assert_equal(ref, f_array_shape())
# shape testing
self._compile_and_check(
[],
unravel_index(indices,
shape_symb,
order=order,
ndim=len(shape)),
[],
UnravelIndex,
)
def check(shape, index_ndim, order):
indices = np.arange(np.product(shape))
# test with scalars and higher-dimensional indices
if index_ndim == 0:
indices = indices[-1]
elif index_ndim == 2:
indices = indices[:, np.newaxis]
indices_symb = theano.shared(indices)
# reference result
ref = np.unravel_index(indices, shape)
def fn(i, d, nd=None):
if nd is None:
return function([], unravel_index(i, d, order=order))
else:
return function([], unravel_index(i, d, order=order, ndim=nd))
# shape given as a tuple
f_array_tuple = fn(indices, shape)
f_symb_tuple = fn(indices_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(indices, 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(indices, shape_symb, len(shape))
np.testing.assert_equal(ref, f_array_symb())
# shape given as a Shape op (unravel_index will use get_vector_length
# to infer the number of dimensions)
indexed_array = theano.shared(np.random.uniform(size=shape_array))
f_array_shape = fn(indices, indexed_array.shape)
np.testing.assert_equal(ref, f_array_shape())
# shape testing
self._compile_and_check([],
unravel_index(indices, shape_symb, order=order, ndim=len(shape)),
[], UnravelIndex)
def fn(i, d, nd=None):
if nd is None:
return function([], unravel_index(i, d, order=order))
else:
return function([],
unravel_index(i, d, order=order, ndim=nd))
def test_unravel_index(self):
def check(shape, index_ndim, order):
indices = np.arange(np.product(shape))
# test with scalars and higher-dimensional indices
if index_ndim == 0:
indices = indices[-1]
elif index_ndim == 2:
indices = indices[:, np.newaxis]
indices_symb = theano.shared(indices)
# reference result
ref = np.unravel_index(indices, shape)
def fn(i, d, nd=None):
if nd is None:
return function([], unravel_index(i, d, order=order))
else:
return function([],
unravel_index(i, d, order=order, ndim=nd))
# shape given as a tuple
f_array_tuple = fn(indices, shape)
f_symb_tuple = fn(indices_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(indices, 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(indices, shape_symb, len(shape))
np.testing.assert_equal(ref, f_array_symb())
# shape given as a Shape op (unravel_index will use get_vector_length
# to infer the number of dimensions)
indexed_array = theano.shared(np.random.uniform(size=shape_array))
f_array_shape = fn(indices, indexed_array.shape)
np.testing.assert_equal(ref, f_array_shape())
# shape testing
self._compile_and_check(
[],
unravel_index(indices,
shape_symb,
order=order,
ndim=len(shape)),
[],
UnravelIndex,
)
for order in ("C", "F"):
for index_ndim in (0, 1, 2):
check((3, ), index_ndim, order)
check((3, 4), index_ndim, order)
check((3, 4, 5), index_ndim, order)
# must specify ndim if length of dims is not fixed
with pytest.raises(ValueError):
unravel_index(theano.tensor.ivector(), theano.tensor.ivector())
# must provide integers
with pytest.raises(TypeError):
unravel_index(theano.tensor.fvector(), (3, 4))
with pytest.raises(TypeError):
unravel_index((3, 4), (3.4, 3.2))
with pytest.raises(ValueError):
unravel_index((3, 4), (3, 3), ndim=5.4)
# dims must be a 1D sequence
with pytest.raises(TypeError):
unravel_index((3, 4), 3)
with pytest.raises(TypeError):
unravel_index((3, 4), ((3, 4), ))
def fn(i, d, nd=None):
if nd is None:
return function([], unravel_index(i, d, order=order))
else:
return function([], unravel_index(i, d, order=order, ndim=nd))
def fn(i, d):
return function([], unravel_index(i, d, order=order))