We can't find the internet
Attempting to reconnect
Something went wrong!
Attempting to reconnect
Top-K Values
Return the k largest values from a vector, ordered from largest to smallest. Select values by their numerical value, not their absolute magnitude.
Your task
Implement top_k(x, k).
xis a one-dimensional tensor of finite values with shape(n,).kis an integer with0 <= k <= n.- Return a tensor of shape
(k,)with x's dtype and device. Leave x unchanged.
For x=[-9, 4, 4, 1] and k=3, return [4, 4, 1]. Both occurrences of 4
count: top-k does not remove duplicates. The value -9 is smaller than 1 even
though its magnitude is larger.
For k=n, return all values in descending order. For k=0, return an empty
vector of shape (0,). For k=1, keep a one-element vector rather than returning
a scalar.
torch.topk(x, k) and jax.lax.top_k(x, k) return both values and their
original indices. Return only the values. Since indices are not part of the
answer, no rule for ordering indices of equal values is needed.
Sorting in descending order and taking the first k values also gives the required result. Do not return the values in their original input order.
Understand the ideas behind this exercise. The Problem tab contains the exact implementation requirements.
Choose matching entries for arithmetic
Elementwise operations act on corresponding entries, after any broadcasting.
For a = [1, 2] and b = [3, 4], (a + b)^2 gives [16, 36]. A dot product
instead multiplies corresponding entries and sums them, giving 11.
Broadcasting can make different shapes compatible. Check that it pairs the intended entries; an operation can run successfully and still compute the wrong thing.
Select the largest values along the intended axis
Top-k selection keeps the largest k entries along an axis. The top two values
of [4, 1, 3], in descending order, are [4, 3]. Their original indices are
[0, 2]. Values and indices are different outputs, so return what the task asks
for.
Check three details: which axis to select along, whether the result must be
sorted, and how to resolve ties. A partial-selection algorithm can find the top
entries without sorting them. Test k = 1 and k equal to the axis length, and
use the task's stated behavior for invalid or empty inputs.
Preserve the requested number type
Mathematically equivalent expressions can give slightly different floating-point results. Integer arithmetic can also overflow before a later conversion to floating point. Follow the Problem tab's requirements for dtype, precision, and output shape.
Your code, saved here
Sign in to access your accepted solutions and saved draft history. You can keep reading the lesson without an account.
Sign in to your accountStuck?
PyTorch reference solution
Sign in to track your attempts toward unlocking the reference solution.