def test_argtopk_nd(self, shp, k_, dtype, sorted, idx_dtype): ndim = len(shp) for axis in range(-ndim, ndim): if isinstance(k_, str): k = eval(k_.replace("n", str(shp[axis]))) else: k = k_ if k == 0: continue x = theano.tensor.tensor( name="x", broadcastable=(False,) * len(shp), dtype=dtype ) y = argtopk(x, k, axis=axis, sorted=sorted, idx_dtype=idx_dtype) fn = theano.function([x], y, mode=self.mode) assert any( [isinstance(n.op, self.op_class) for n in fn.maker.fgraph.apply_nodes] ) size = reduce(int.__mul__, shp) xval = gen_unique_vector(size, dtype).reshape(shp) yval = fn(xval) idx = slice(-k, None) if k > 0 else slice(-k) l = axis % ndim r = ndim - l idx = (slice(None),) * l + (idx,) + (slice(None),) * (r - 1) goal = np.argsort(xval, axis=axis)[idx].astype(idx_dtype) assert np.all(np.sort(yval, axis=axis) == np.sort(goal, axis=axis))
def test_argtopk_nd(self, shp, k_, dtype, sorted, idx_dtype): ndim = len(shp) for axis in range(-ndim, ndim): if isinstance(k_, str): k = eval(k_.replace('n', str(shp[axis]))) else: k = k_ if k == 0: continue x = theano.tensor.tensor( name='x', broadcastable=(False,) * len(shp), dtype=dtype) y = argtopk(x, k, axis=axis, sorted=sorted, idx_dtype=idx_dtype) fn = theano.function([x], y, mode=self.mode) assert any([isinstance(n.op, self.op_class) for n in fn.maker.fgraph.apply_nodes]) size = reduce(int.__mul__, shp) xval = gen_unique_vector(size, dtype).reshape(shp) yval = fn(xval) idx = slice(-k, None) if k > 0 else slice(-k) l = axis % ndim r = ndim - l idx = (slice(None),) * l + (idx,) + (slice(None),) * (r - 1) goal = np.argsort(xval, axis=axis)[idx].astype(idx_dtype) assert np.all(np.sort(yval, axis=axis) == np.sort(goal, axis=axis))
def test_argtopk_sanity(self, dtype, idx_dtype, axis, sorted): x = tensor.vector(name='x', dtype=dtype) fn = theano.function([x], argtopk(x, 1, axis=axis, sorted=sorted, idx_dtype=idx_dtype)) xval = np.asarray([1]).astype(dtype) yval = fn(xval) assert yval == np.asarray([0], dtype=idx_dtype) assert yval.dtype == np.dtype(idx_dtype)
def test_argtopk_sanity(self, dtype, idx_dtype, axis, sorted): x = tensor.vector(name='x', dtype=dtype) fn = theano.function([x], argtopk(x, 1, axis=axis, sorted=sorted, idx_dtype=idx_dtype), mode=self.mode) assert any([isinstance(n.op, self.op_class) for n in fn.maker.fgraph.apply_nodes]) xval = np.asarray([1]).astype(dtype) yval = fn(xval) assert yval == np.asarray([0], dtype=idx_dtype) assert yval.dtype == np.dtype(idx_dtype)
def test_argtopk_sanity(self, dtype, idx_dtype, axis, sorted): x = tensor.vector(name="x", dtype=dtype) fn = theano.function( [x], argtopk(x, 1, axis=axis, sorted=sorted, idx_dtype=idx_dtype), mode=self.mode, ) assert any( [isinstance(n.op, self.op_class) for n in fn.maker.fgraph.apply_nodes] ) xval = np.asarray([1]).astype(dtype) yval = fn(xval) assert yval == np.asarray([0], dtype=idx_dtype) assert yval.dtype == np.dtype(idx_dtype)
def test_argtopk_1d_collision(self, size, k, dtype, sorted): # with non-unique kth max value if isinstance(k, str): k = eval(k.replace('n', str(size))) x = theano.tensor.vector(name='x', dtype=dtype) y = argtopk(x, k, sorted=sorted, idx_dtype='int32') fn = theano.function([x], y) xval = np.repeat(np.random.uniform(-100., 100., size=size // 2).astype(dtype), 2) xval = xval[np.random.permutation(size)] yval = fn(xval) idx = (slice(-k, None) if k > 0 else slice(-k)) goal = np.argsort(xval)[idx].astype('int32') utt.assert_allclose(np.sort(xval[yval]), np.sort(xval[goal]))
def test_argtopk_1d_collision(self, size, k, dtype, sorted): # with non-unique kth max value if isinstance(k, str): k = eval(k.replace('n', str(size))) x = theano.tensor.vector(name='x', dtype=dtype) y = argtopk(x, k, sorted=sorted, idx_dtype='int32') fn = theano.function([x], y) xval = np.repeat( np.random.uniform(-100., 100., size=size // 2).astype(dtype), 2) xval = xval[np.random.permutation(size)] yval = fn(xval) idx = (slice(-k, None) if k > 0 else slice(-k)) goal = np.argsort(xval)[idx].astype('int32') utt.assert_allclose(np.sort(xval[yval]), np.sort(xval[goal]))
def test_argtopk_1d(self, size, k, dtype, sorted, idx_dtype): if isinstance(k, str): k = eval(k.replace('n', str(size))) x = theano.tensor.vector(name='x', dtype=dtype) y = argtopk(x, k, sorted=sorted, idx_dtype=idx_dtype) fn = theano.function([x], y) # assert local_useless_topk opt is done properly assert 1 == len(fn.maker.fgraph.outputs[0].owner.outputs) # generate a all-unique array xval = gen_unique_vector(size, dtype) yval = fn(xval) idx = (slice(-k, None) if k > 0 else slice(-k)) goal = np.argsort(xval)[idx].astype(idx_dtype) # due to uniqueness, we expect indices same assert np.all(xval[np.sort(yval)] == xval[np.sort(goal)])
def test_argtopk_1d_collision(self, size, k, dtype, sorted): # with non-unique kth max value if isinstance(k, str): k = eval(k.replace('n', str(size))) x = theano.tensor.vector(name='x', dtype=dtype) y = argtopk(x, k, sorted=sorted, idx_dtype='int32') # DebugMode won't like the index change on collision on CPU # So don't use DebugMode here. mode = self.mode if isinstance(self.mode, theano.compile.DebugMode): mode = theano.Mode(optimizer=mode.optimizer) fn = theano.function([x], y, mode=mode) assert any([isinstance(n.op, self.op_class) for n in fn.maker.fgraph.apply_nodes]) xval = np.repeat(np.random.uniform(-100., 100., size=size // 2).astype(dtype), 2) xval = xval[np.random.permutation(size)] yval = fn(xval) idx = (slice(-k, None) if k > 0 else slice(-k)) goal = np.argsort(xval)[idx].astype('int32') utt.assert_allclose(np.sort(xval[yval]), np.sort(xval[goal]))
def test_argtopk_1d_collision(self, size, k, dtype, sorted): # with non-unique kth max value if isinstance(k, str): k = eval(k.replace("n", str(size))) x = theano.tensor.vector(name="x", dtype=dtype) y = argtopk(x, k, sorted=sorted, idx_dtype="int32") # DebugMode won't like the index change on collision on CPU # So don't use DebugMode here. mode = self.mode if isinstance(self.mode, theano.compile.DebugMode): mode = theano.Mode(optimizer=mode.optimizer) fn = theano.function([x], y, mode=mode) assert any( [isinstance(n.op, self.op_class) for n in fn.maker.fgraph.apply_nodes] ) xval = np.repeat( np.random.uniform(-100.0, 100.0, size=size // 2).astype(dtype), 2 ) xval = xval[np.random.permutation(size)] yval = fn(xval) idx = slice(-k, None) if k > 0 else slice(-k) goal = np.argsort(xval)[idx].astype("int32") utt.assert_allclose(np.sort(xval[yval]), np.sort(xval[goal]))