tqchen commented on code in PR #16351:
URL: https://github.com/apache/tvm/pull/16351#discussion_r1445025872
##########
src/relax/op/tensor/sort.h:
##########
@@ -43,6 +43,28 @@ namespace relax {
*/
Expr sort(Expr data, int axis, bool descending);
+/*!
+ * \brief Performs sorting along the given axis and returns an array of
indices.
+ * \param data The input tensor.
+ * \param axis The axis to sort on.
+ * \param descending Whether to sort in descending order.
+ * \param dtype The data type of the output indices.
+ * \return The computed result.
+ */
+Expr argsort(Expr data, int axis, bool descending, DataType dtype);
+
+/*!
+ * \brief Get the top k elements in an input tensor along the given axis.
+ * \param data The input tensor.
+ * \param k Number of top elements.
+ * \param axis The axis to sort on.
+ * \param ret_type The return type, can be set to one of [both, values,
indices].
+ * \param largest Whether to return largest or smallest elements.
+ * \param dtype The data type of the indices output.
Review Comment:
rename to sorting.h to be consistent with
https://data-apis.org/array-api/latest/API_specification/index.html
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
