aesara.tensor.argtopk

Contents

aesara.tensor.argtopk#

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

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.