def test_combined_sanity(self, dtype, idx_dtype, axis, sorted):
x = vector(name="x", dtype=dtype)
yv, yi = topk_and_argtopk(x,
1,
axis=axis,
sorted=sorted,
idx_dtype=idx_dtype)
fn = aesara.function([x], [yv, yi], 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)
yvval, yival = fn(xval)
assert yival == np.asarray([0], dtype=idx_dtype)
utt.assert_allclose(xval, yvval)
assert yvval.dtype == xval.dtype
assert yival.dtype == np.dtype(idx_dtype)
def test_combined_1d(self, size, k, dtype, sorted, idx_dtype):
if isinstance(k, str):
k = eval(k.replace("n", str(size)))
x = vector(name="x", dtype=dtype)
yv, yi = topk_and_argtopk(x, k, sorted=sorted, idx_dtype=idx_dtype)
fn = aesara.function([x], [yv, yi], mode=self.mode)
assert any(
isinstance(n.op, self.op_class)
for n in fn.maker.fgraph.apply_nodes)
# generate a all-unique array
xval = gen_unique_vector(size, dtype)
yvval, yival = fn(xval)
idx = slice(-k, None) if k > 0 else slice(-k)
goali = np.argsort(xval)[idx].astype(idx_dtype)
goalv = xval[goali]
# due to uniqueness, we expect indices same
assert np.all(xval[np.sort(yival)] == xval[np.sort(goali)])
utt.assert_allclose(np.sort(yvval), goalv)
def test_combined_infer_shape(self, shp, k_):
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 = tensor(name="x",
broadcastable=(False, ) * len(shp),
dtype=aesara.config.floatX)
yv, yi = topk_and_argtopk(x,
k,
axis=axis,
sorted=False,
idx_dtype="int32")
size = reduce(int.__mul__, shp)
xval = gen_unique_vector(size, aesara.config.floatX).reshape(shp)
self._compile_and_check([x], [yv, yi], [xval], TopKOp)
