mxnet.npx.topk

topk(data, axis=-1, k=1, ret_typ='indices', is_ascend=False, dtype='float32')
Returns the indices of the top k elements in an input array along the given

axis (by default). If ret_type is set to ‘value’ returns the value of top k elements (instead of indices). In case of ret_type = ‘both’, both value and index would be returned. The returned elements will be sorted.

Parameters
  • data (NDArray) – The input array

  • axis (int or None, optional, default='-1') – Axis along which to choose the top k indices. If not given, the flattened array is used. Default is -1.

  • k (int, optional, default='1') – Number of top elements to select, should be always smaller than or equal to the element number in the given axis. A global sort is performed if set k < 1.

  • ret_typ ({'both', 'indices', 'mask', 'value'},optional, default='indices') – The return type. “value” means to return the top k values, “indices” means to return the indices of the top k values, “mask” means to return a mask array containing 0 and 1. 1 means the top k values. “both” means to return a list of both values and indices of top k elements.

  • is_ascend (boolean, optional, default=0) – Whether to choose k largest or k smallest elements. Top K largest elements will be chosen if set to false.

  • dtype ({'float16', 'float32', 'float64', 'int32', 'int64', 'uint8'},) – optional, default=’float32’ DType of the output indices when ret_typ is “indices” or “both”. An error will be raised if the selected data type cannot precisely represent the indices.

Returns

out – The output of this function.

Return type

NDArray or list of NDArrays

Example

>>> x = np.array([[0.3, 0.2, 0.4], [0.1, 0.3, 0.2]])

returns an index of the largest element on last axis

>>> npx.topk(x)
array([[2.],
       [1.]])

returns the value of top-2 largest elements on last axis

>>> npx.topk(x, ret_typ='value', k=2)
array([[0.4, 0.3],
       [0.3, 0.2]])

returns the value of top-2 smallest elements on last axis

>>> npx.topk(x, ret_typ='value', k=2, is_ascend=1)
array([[0.2, 0.3],
       [0.1, 0.2]])

returns the value of top-2 largest elements on axis 0

>>> npx.topk(x, axis=0, ret_typ='value', k=2)
array([[0.3, 0.3, 0.4],
       [0.1, 0.2, 0.2]])

flattens and then returns list of both values and indices

>>> npx.topk(x, ret_typ='both', k=2)
[array([[0.4, 0.3], [0.3, 0.2]]),
 array([[2., 0.], [1., 2.]])]