wehub-resource-sync
129 lines
4.1 KiB
Python
129 lines
4.1 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.
|
|
# pylint: disable=invalid-name
|
|
"""Search operators."""
|
|
|
|
from ..expr import Expr
|
|
from . import _ffi_api
|
|
|
|
|
|
def where(condition: Expr, x1: Expr, x2: Expr) -> Expr:
|
|
"""Selecting elements from either the input tensors depending on the value of the
|
|
condition.
|
|
|
|
For a given position, return the corresponding value in `x1` if `condition` is True,
|
|
and return the corresponding value in `x2` otherwise.
|
|
|
|
Parameters
|
|
----------
|
|
condition : relax.Expr
|
|
When True, yield `x1`; otherwise, yield `x2`.
|
|
Must be broadcasting compatible with `x1` and `x2`.
|
|
Must have boolean dtype.
|
|
|
|
x1 : relax.Expr
|
|
The first input tensor.
|
|
Must be broadcasting compatible with `condition` and `x2`.
|
|
|
|
x2 : relax.Expr
|
|
The second input tensor.
|
|
Must be broadcasting compatible with `condition` and `x1`.
|
|
|
|
Returns
|
|
-------
|
|
result : relax.Expr
|
|
The result tensor.
|
|
"""
|
|
return _ffi_api.where(condition, x1, x2) # type: ignore
|
|
|
|
|
|
def argmax(x: Expr, axis: int | None = None, keepdims: bool = False) -> Expr:
|
|
"""Computes the argmax of tensor elements over given axis.
|
|
|
|
Parameters
|
|
----------
|
|
x : relax.Expr
|
|
The input data tensor
|
|
|
|
axis : Optional[int]
|
|
Axis along which an argmax operation is performed.
|
|
The default, axis=None, will compute the argmax of all elements in the input tensor.
|
|
Negative indexing is supported.
|
|
|
|
keepdims : bool
|
|
If this is set to True, the axis being reduced is left in the result as dimensions
|
|
with size one.
|
|
With this option, the result will broadcast correctly against the input tensor.
|
|
|
|
Returns
|
|
-------
|
|
result : relax.Expr
|
|
The computed result.
|
|
"""
|
|
return _ffi_api.argmax(x, axis, keepdims) # type: ignore
|
|
|
|
|
|
def argmin(x: Expr, axis: int | None = None, keepdims: bool = False) -> Expr:
|
|
"""Computes the argmin of tensor elements over given axis.
|
|
|
|
Parameters
|
|
----------
|
|
x : relax.Expr
|
|
The input data tensor
|
|
|
|
axis : Optional[int]
|
|
Axis along which an argmin operation is performed.
|
|
The default, axis=None, will compute the argmin of all elements in the
|
|
input tensor. Negative indexing is supported.
|
|
|
|
keepdims : bool
|
|
If this is set to True, the axis being reduced is left in the result as
|
|
dimensions with size one.
|
|
With this option, the result will broadcast correctly against the input tensor.
|
|
|
|
Returns
|
|
-------
|
|
result : relax.Expr
|
|
The computed result.
|
|
"""
|
|
return _ffi_api.argmin(x, axis, keepdims) # type: ignore
|
|
|
|
|
|
def bucketize(input_tensor, boundaries, out_int32=False, right=False):
|
|
"""Returns the indices of the buckets to which each value in the input belongs.
|
|
|
|
Parameters
|
|
----------
|
|
input_tensor : relax.Expr
|
|
N-D tensor containing the search values.
|
|
|
|
boundaries : relax.Expr
|
|
1-D tensor, must contain a strictly increasing sequence, or the return value is undefined.
|
|
|
|
out_int32 : Optional[bool]
|
|
Indicate the output data type. int32 if True, int64 otherwise. Default=False
|
|
|
|
right : Optional[bool]
|
|
Determines the behavior for values in boundaries. Similar to torch.bucketize
|
|
|
|
Returns
|
|
-------
|
|
result : relax.Expr
|
|
The computed result with same shape as input_tensor.
|
|
"""
|
|
return _ffi_api.bucketize(input_tensor, boundaries, out_int32, right)
|