wehub-resource-sync
118 lines
3.3 KiB
Python
118 lines
3.3 KiB
Python
# Licensed to the Apache Software Foundation (ASF) under one
|
|
# or more contributor license agreements. See the NOTICE file
|
|
# distributed with this work for additional information
|
|
# regarding copyright ownership. The ASF licenses this file
|
|
# to you under the Apache License, Version 2.0 (the
|
|
# "License"); you may not use this file except in compliance
|
|
# with the License. You may obtain a copy of the License at
|
|
#
|
|
# http://www.apache.org/licenses/LICENSE-2.0
|
|
#
|
|
# Unless required by applicable law or agreed to in writing,
|
|
# software distributed under the License is distributed on an
|
|
# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
|
|
# KIND, either express or implied. See the License for the
|
|
# specific language governing permissions and limitations
|
|
# under the License.
|
|
"""Sortings operators."""
|
|
|
|
from ..expr import Constant, Expr
|
|
from . import _ffi_api
|
|
|
|
|
|
def sort(x: Expr, axis: int = -1, descending: bool = False):
|
|
"""Performs sorting along the given axis and returns an array
|
|
in sorted order.
|
|
|
|
Parameters
|
|
----------
|
|
x : relax.Expr
|
|
The input tensor.
|
|
|
|
axis : int
|
|
Axis along which to sort the input tensor.
|
|
By default the last axis of the input is used.
|
|
|
|
descending : bool
|
|
Whether to sort in descending order, the default is False
|
|
|
|
Returns
|
|
-------
|
|
out : relax.Expr
|
|
Sorted tensor.
|
|
|
|
"""
|
|
return _ffi_api.sort(x, axis, descending) # type: ignore
|
|
|
|
|
|
def argsort(data: Expr, axis: int = -1, descending: bool = False, dtype: str = "int32"):
|
|
"""Performs sorting along the given axis and returns an array of indices
|
|
having same shape as an input array that index data in sorted order.
|
|
|
|
Parameters
|
|
----------
|
|
data : relax.Expr
|
|
The input data tensor.
|
|
|
|
axis : int
|
|
Axis long which to sort the input tensor.
|
|
|
|
descending : bool
|
|
Whether to sort in descending order, the default is False
|
|
|
|
dtype : str
|
|
The data type of the output indices.
|
|
|
|
Returns
|
|
-------
|
|
out : relax.Expr
|
|
Tensor with same shape as data.
|
|
"""
|
|
return _ffi_api.argsort(data, axis, descending, dtype) # type: ignore
|
|
|
|
|
|
def topk(
|
|
data: Expr,
|
|
k: int = 1,
|
|
axis: int = -1,
|
|
ret_type: str = "both",
|
|
largest: bool = True,
|
|
dtype: str = "int32",
|
|
):
|
|
"""Get the top k elements in an input tensor along the given axis.
|
|
|
|
ret_type specifies the return type, can be one of ("both", "values", "indices").
|
|
|
|
Parameters
|
|
----------
|
|
data : relax.Expr
|
|
The input data tensor.
|
|
|
|
k : int
|
|
Number of top elements to select. Return all elements if k < 1.
|
|
|
|
axis : int
|
|
Axis long which to sort the input tensor.
|
|
|
|
ret_type: str
|
|
The return type [both, values, indices].
|
|
"both": return both top k data and indices.
|
|
"values": return top k data only.
|
|
"indices": return top k indices only.
|
|
|
|
largest : bool
|
|
Whether to return largest or smallest elements.
|
|
The k smallest elements are returned if largest is False.
|
|
|
|
dtype : str
|
|
The data type of the indices output.
|
|
|
|
Returns
|
|
-------
|
|
out : relax.Expr or List[relax.Expr]
|
|
The computed result.
|
|
"""
|
|
if isinstance(k, Constant):
|
|
k = k.data.numpy().item()
|
|
return _ffi_api.topk(data, k, axis, ret_type, largest, dtype) # type: ignore
|