topK(_:)

Declaration

func topK(_ k: Int) -> (values: MLTensor, indices: MLTensor)

Parameters

• k:

  The number of largest values to return.

Return Value

A tuple of (values, indices) with values containing the top k values along the last axis and indices, a Int32 tensor, containing the indices to the corresponding values.

Discussion

For example:

let x = MLTensor(shape: [1, 5], scalars: [1.0, 2.0, 3.0, 4.0, 5.0])
let (values, indices) = x.topK(3)
// values = [[5.0, 4.0, 3.0]]
// indices = [[4, 3, 2]]