Contents

topK(_:axis:findLargest:)

Adds a top-k operation to the current graph.

Declaration

func topK(_ k: Int, axis: Int, findLargest: Bool) -> (values: BNNSGraph.Builder.Tensor<T>, indices: BNNSGraph.Builder.Tensor<Int32>)

Discussion

The output of this graph is a pair of tensors (values, indices) that contains the top-k values and indices from the tensor from the given dimension. If largest is true, then output the largest k values and their corresponding indices, otherwise output the smallest.