Add `topK` operator (a sliced `sort`)

Start by reviewing the proposed MLGraphBuilder signatures and MLOpSupportLimits entries in the issue, then compare the referenced Apple, ONNX, TFLite/LiteRT, and PyTorch APIs. Decide the topK and sort definitions, including options, return shape, tie behavior, NaN handling, and backend support; done means the API proposal resolves these choices and specifies behavior that can be implemented consistently.

Forked from https://github.com/webmachinelearning/webnn/issues/573#issuecomment-5027913913.

@mtavenrath, 👀 for topK (a component of an MoE graph), I see decent backend support that WebNN could use:

• Apple tensor_operation.topk
• ONNX TopK
• TFLite/LiteRT - tfl.topk_v2
• PyTorch torch.topk

Also, for an even more core operator, it seems a sort operator should exist, like StableHLO's sort (interestingly StableHLO has sort, but not topK). Since many ML libraries lack sort, sorting is evidently either unused by models directly, or models are achieving sort by abusing topK with full k. So adding sort would be more for foundational composition reasons. Depending on the backend, topK can be composed from sort and slice, whereas sort can be easily composed from topK and (k == full dimension size).

For MoE itself, that feels more like a mini-graph rather than an operator (granted, we currently have gru/lstm which are mini-graphs), and I'm seeing different definitions of MoE across libraries, suggesting some instability still in the field 🤔. None of the primary Chromium ML backends have it. Though ONNX Runtime has a contributor operator. In any case, a particular variant of MoE should be composable, given topK and existing WebNN operators.

Prototype

Possible definition:

dictionary MLTopKOptions : MLOperatorOptions
{
    unsigned long axis;
    boolean ascending = false; // Default to descending order.
    boolean sorted = true;
};

dictionary MLSortOptions : MLOperatorOptions
{
    unsigned long axis;
    boolean ascending = true; // Default to ascending order.
};
// Maybe these dictionaries could be combined cleverly?

partial interface MLGraphBuilder
{
    sequence<MLOperand> topK(MLOperand input, unsigned long k, optional MLTopKOptions options = {});
    // ⚠️ Would returning a named struct be clearer?
    //     Like: dictionary MLTopKResult {MLOperand values; MLOperand indices}?
    //     Or is returning a sequence simpler for existing logic given split/gru/lstm?

    sequence<MLOperand> sort(MLOperand input, optional MLSortOptions options = {});
};

partial dictionary MLOpSupportLimits
{
    MLSingleInputSupportLimits topK;
    MLSingleInputSupportLimits sort;
};

Considerations

• TFLite always sorts the topK output indices. So passing sorted as false would make no difference - would still be true.
• TFLite always return descending order, and so for ascending order, you'd have to transpose the output, and then slice it to k.
• TFLite is missing an axis parameter, always sorting the last dimension. So you'd have to transpose the input first.
• I didn't include any "stable sorting" flag like StableHLO, since some API's break ties based on the element index location (which essentially means stable sorting is always enabled), and others don't guarantee it (meaning essentially false).
• NaN's are an annoying conformance topic 🤔 (see here).

Full Reference of other API's

Apple MIL tensor_operation.topk

classcoremltools.converters.mil.mil.ops.defs.iOS15.tensor_operation.topk(**kwargs)

Returns a tensor containing top or bottom k values and the corresponding indices of the input tensor along a given axis.

Parameters:

• x: <*?, T> (Required) - Input tensor.
• k: const (Optional) - Defaults to 1. Number of values/indices to be computed along each axis.
• axis: const (Optional) - Defaults to -1 (last dimension). Axis to perform the operation.
• ascending: const (Optional) - Defaults to False, sort in descending order. True to sort in ascending order.

Attributes:

• T: fp16, fp32, int32

Returns:

• tensor<*?, T>: Values of top/bottom k elements.
• tensor<*?, int32>: Indices of the top/bottom k elements along axis.

ONNX TopK

Retrieve the top-K largest or smallest elements along a specified axis. Given an input tensor of shape [a_0, a_1, …, a_{n-1}] and integer argument k, return two outputs:

• Value tensor of shape [a_0, a_1, …, a_{axis-1}, k, a_{axis+1}, … a_{n-1}] which contains the values of the top k elements along the specified axis
• Index tensor of shape [a_0, a_1, …, a_{axis-1}, k, a_{axis+1}, … a_{n-1}] which contains the indices of the top k elements (original indices from the input tensor).
• If “largest” is 1 (the default value) then the k largest elements are returned.
• If “sorted” is 1 (the default value) then the resulting k elements will be sorted.
• If “sorted” is 0, order of returned ‘Values’ and ‘Indices’ are undefined.

Given two equivalent values, this operator uses the indices along the axis as a tiebreaker. That is, the element with the lower index will appear first.

Attributes:

• axis - INT (default is -1): Dimension on which to do the sort. Negative value means counting dimensions from the back. Accepted range is [-r, r-1] where r = rank(input).
• largest - INT (default is 1): Whether to return the top-K largest or smallest elements.
• sorted - INT (default is 1): Whether to return the elements in sorted order.

Inputs:

• X (heterogeneous) - T: Tensor of shape [a_0, a_1, …, a_{n-1}]
• K (heterogeneous) - tensor(int64): A 1-D tensor containing a single positive value corresponding to the number of top elements to retrieve

Outputs:

Values (heterogeneous) - T: Tensor of shape [a_0, a_1, …, a_{axis-1}, k, a_{axis+1}, … a_{n-1}] containing top K values from the input tensor
Indices (heterogeneous) - I: Tensor of shape [a_0, a_1, …, a_{axis-1}, k, a_{axis+1}, … a_{n-1}] containing the corresponding input tensor indices for the top K values.

TFLite/LiteRT tfl.topk_v2 (TFL::TopKV2Op)

Returns the top k largest element along each last dimensional slice of input and the indices of values within the last dimension of the input tensor. Results are always sorted in the descending order.

Operands:

• input: tensor of 32-bit float or 8-bit signless integer or 16-bit signless integer or 32-bit signless integer or 64-bit signless integer or 8-bit unsigned integer or QI8 type or QUI8 type values
• k: tensor of 16-bit signless integer or 32-bit signless integer values
  Results:
• values: tensor of 32-bit float or 8-bit signless integer or 16-bit signless integer or 32-bit signless integer or 64-bit signless integer or 8-bit unsigned integer or QI8 type or QUI8 type values
• indices: tensor of 16-bit signless integer or 32-bit signless integer values

PyTorch torch.topk

torch.topk(input, k, dim=None, largest=True, sorted=True, *, out=None)

Returns the k largest elements of the given input tensor along a given dimension.

• If dim is not given, the last dimension of the input is chosen.
• If largest is False then the k smallest elements are returned.
• A namedtuple of (values, indices) is returned with the values and indices of the largest k elements of each row of the input tensor in the given dimension dim.
• The boolean option sorted if True, will make sure that the returned k elements are themselves sorted

When using torch.topk, the indices of tied elements are not guaranteed to be stable and may vary across different invocations.

Parameters:

• input (Tensor) – the input tensor.
• k (int) – the k in “top-k”
• dim (int, optional) – the dimension to sort along
• largest (bool, optional) – controls whether to return largest or smallest elements
• sorted (bool, optional) – controls whether to return the elements in sorted order

Keyword Arguments:

• out (tuple, optional) – the output tuple of (Tensor, LongTensor) that can be optionally given to be used as output buffers