aesara.tensor.topk

Contents

aesara.tensor.topk#

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

Returns the 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.

  • axis (integer or None) – Upon which axis shall the operation be performed on. If None, works on flattened array.

  • sorted (bool) –

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

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

  • idx_dtype (string) – Specify output dtype used in indices, defaults to int64, must be integer type. This option is here because indices are needed for gradient.

Return type:

Tensor variable with same dtype as x.

Notes

  • sorted=True is not supported yet.