topk_op

→vector<TensorImplPtr>

int topk_op(const int & a, int k, int axis)

source edit

Python wrappers

pylucid.ops.topkfunction

Top-k values and their indices along an axis.

Selects the k largest elements along axis and returns them together with their positions in the original tensor.

See Also

sort_op : Full sort (no truncation). argsort_op : Sort indices only.

Parameters

aTensorImplPtr

Input tensor.

kint64

Number of top elements to return. Must satisfy 0 < k <= a.shape[axis].

axisint

Axis along which to select.

Returns

vector<TensorImplPtr>

[values, indices] where both have shape equal to a.shape with axis reduced to k. values has the dtype of a; indices has dtype I32.

Notes

values is differentiable for F32 / F64 inputs via IndexScatterBackward, which scatters the incoming gradient back to the selected positions in the original tensor. indices is never differentiable.