aesara.tensor.argtopk

aesara.tensor.argtopk(x, kth, axis=-1, sorted=True, idx_dtype='int64')

Returns the indices of k-largest elements along an axis.

Parameters:

• x (tensor instance) –

• kth (integer constant/variable) – Must not be 0. If negative, gives k-smallest elements instead.

• sorted (bool) –

  NOTE: NOT IMPLEMENTED YET, USE False FOR NOW. Defaults to True

  If True, the result array of corresponding indices would be sorted in descending order.

axis: integer, tuple/list of integers, or None

Upon which axis shall the operation be performed on. If None, works on flattened array.

idx_dtype: string

Specify output dtype, defaults to int64, must be integer type.

Return type:

Tensor variable with dtype specified in idx_dtype.

Notes

• sorted=True is not supported yet.

• If the top-k-th value is not unique, we cannot guarantee the output indices are deterministically chosen.