chore: import upstream snapshot with attribution
This commit is contained in:
wehub-resource-sync
@@ -0,0 +1,484 @@
|
||||
# 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=redefined-builtin, invalid-name
|
||||
"""Relax binary arithmetic and comparison operators."""
|
||||
|
||||
from ..expr import Expr
|
||||
from . import _ffi_api
|
||||
|
||||
###################### Arithmetic operators ######################
|
||||
|
||||
|
||||
def add(x1: Expr, x2: Expr) -> Expr:
|
||||
"""Addition with numpy-style broadcasting.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
x1 : Expr
|
||||
The first input tensor.
|
||||
x2 : Expr
|
||||
The second input tensor.
|
||||
|
||||
Returns
|
||||
-------
|
||||
result : Expr
|
||||
The computed result.
|
||||
|
||||
Examples
|
||||
--------
|
||||
.. code:: python
|
||||
|
||||
bb = relax.BlockBuilder()
|
||||
a = relax.Var("a", relax.TensorType(shape=(2, 3), dtype="float32"))
|
||||
b = relax.Var("b", relax.TensorType(shape=(2, 1), dtype="float32"))
|
||||
c = bb.normalize(relax.op.add(a, b)) # c has TensorType(shape=(2, 3), dtype="float32")
|
||||
"""
|
||||
return _ffi_api.add(x1, x2) # type: ignore
|
||||
|
||||
|
||||
def divide(x1: Expr, x2: Expr) -> Expr:
|
||||
"""Division with numpy-style broadcasting.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
x1 : relax.Expr
|
||||
The first input tensor.
|
||||
x2 : relax.Expr
|
||||
The second input tensor.
|
||||
|
||||
Returns
|
||||
-------
|
||||
result : relax.Expr
|
||||
The computed result.
|
||||
"""
|
||||
return _ffi_api.divide(x1, x2) # type: ignore
|
||||
|
||||
|
||||
def floor_divide(x1: Expr, x2: Expr) -> Expr:
|
||||
"""Floor division with numpy-style broadcasting.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
x1 : relax.Expr
|
||||
The first input tensor.
|
||||
x2 : relax.Expr
|
||||
The second input tensor.
|
||||
|
||||
Returns
|
||||
-------
|
||||
result : relax.Expr
|
||||
The computed result.
|
||||
"""
|
||||
return _ffi_api.floor_divide(x1, x2) # type: ignore
|
||||
|
||||
|
||||
def log_add_exp(x1: Expr, x2: Expr) -> Expr:
|
||||
"""
|
||||
Compute the log of the sum of exponentials of the inputs, element-wise.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
x1 : Expr
|
||||
The first input tensor.
|
||||
x2 : Expr
|
||||
The second input tensor.
|
||||
|
||||
Returns
|
||||
-------
|
||||
Expr
|
||||
The element-wise log-sum-exp of `x1` and `x2`.
|
||||
"""
|
||||
return _ffi_api.log_add_exp(x1, x2)
|
||||
|
||||
|
||||
def multiply(x1: Expr, x2: Expr) -> Expr:
|
||||
"""Multiplication with numpy-style broadcasting.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
x1 : Expr
|
||||
The first input tensor.
|
||||
x2 : Expr
|
||||
The second input tensor.
|
||||
|
||||
Returns
|
||||
-------
|
||||
result : Expr
|
||||
The computed result.
|
||||
"""
|
||||
return _ffi_api.multiply(x1, x2) # type: ignore
|
||||
|
||||
|
||||
def power(x1: Expr, x2: Expr):
|
||||
"""Power with numpy-style broadcasting.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
x1 : relax.Expr
|
||||
The first input tensor.
|
||||
x2 : relax.Expr
|
||||
The second input tensor.
|
||||
|
||||
Returns
|
||||
-------
|
||||
result : relax.Expr
|
||||
The computed result.
|
||||
"""
|
||||
return _ffi_api.power(x1, x2) # type: ignore
|
||||
|
||||
|
||||
def atan2(x1: Expr, x2: Expr) -> Expr:
|
||||
"""Atan2 with numpy-style broadcasting.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
x1 : relax.Expr
|
||||
The first input tensor (y-coordinates).
|
||||
x2 : relax.Expr
|
||||
The second input tensor (x-coordinates).
|
||||
|
||||
Returns
|
||||
-------
|
||||
result : relax.Expr
|
||||
The computed result.
|
||||
"""
|
||||
return _ffi_api.atan2(x1, x2) # type: ignore
|
||||
|
||||
|
||||
def subtract(x1: Expr, x2: Expr) -> Expr:
|
||||
"""Subtraction with numpy-style broadcasting.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
x1 : relax.Expr
|
||||
The first input tensor.
|
||||
x2 : relax.Expr
|
||||
The second input tensor.
|
||||
|
||||
Returns
|
||||
-------
|
||||
result : relax.Expr
|
||||
The computed result.
|
||||
"""
|
||||
return _ffi_api.subtract(x1, x2) # type: ignore
|
||||
|
||||
|
||||
def mod(x1: Expr, x2: Expr) -> Expr:
|
||||
"""Modulo with numpy-style broadcasting.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
x1 : Expr
|
||||
The first input tensor.
|
||||
x2 : Expr
|
||||
The second input tensor.
|
||||
"""
|
||||
return _ffi_api.mod(x1, x2) # type: ignore
|
||||
|
||||
|
||||
def floor_mod(x1: Expr, x2: Expr) -> Expr:
|
||||
"""Floor modulo with numpy-style broadcasting.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
x1 : Expr
|
||||
The first input tensor.
|
||||
x2 : Expr
|
||||
The second input tensor.
|
||||
"""
|
||||
return _ffi_api.floor_mod(x1, x2) # type: ignore
|
||||
|
||||
|
||||
###################### Comparison operators ######################
|
||||
|
||||
|
||||
def equal(x1: Expr, x2: Expr) -> Expr:
|
||||
"""Broadcasted element-wise test for (lhs == rhs).
|
||||
|
||||
Parameters
|
||||
----------
|
||||
x1 : relax.Expr
|
||||
The first input tensor.
|
||||
x2 : relax.Expr
|
||||
The second input tensor.
|
||||
|
||||
Returns
|
||||
-------
|
||||
result : relax.Expr
|
||||
The computed result.
|
||||
"""
|
||||
return _ffi_api.equal(x1, x2) # type: ignore
|
||||
|
||||
|
||||
def greater(x1: Expr, x2: Expr) -> Expr:
|
||||
"""Broadcasted element-wise test for (lhs > rhs).
|
||||
|
||||
Parameters
|
||||
----------
|
||||
x1 : relax.Expr
|
||||
The first input tensor.
|
||||
x2 : relax.Expr
|
||||
The second input tensor.
|
||||
|
||||
Returns
|
||||
-------
|
||||
result : relax.Expr
|
||||
The computed result.
|
||||
"""
|
||||
return _ffi_api.greater(x1, x2) # type: ignore
|
||||
|
||||
|
||||
def greater_equal(x1: Expr, x2: Expr) -> Expr:
|
||||
"""Broadcasted element-wise test for (lhs >= rhs).
|
||||
|
||||
Parameters
|
||||
----------
|
||||
x1 : relax.Expr
|
||||
The first input tensor.
|
||||
x2 : relax.Expr
|
||||
The second input tensor.
|
||||
|
||||
Returns
|
||||
-------
|
||||
result : relax.Expr
|
||||
The computed result.
|
||||
"""
|
||||
return _ffi_api.greater_equal(x1, x2) # type: ignore
|
||||
|
||||
|
||||
def less(x1: Expr, x2: Expr) -> Expr:
|
||||
"""Broadcasted element-wise test for (lhs < rhs).
|
||||
|
||||
Parameters
|
||||
----------
|
||||
x1 : relax.Expr
|
||||
The first input tensor.
|
||||
x2 : relax.Expr
|
||||
The second input tensor.
|
||||
|
||||
Returns
|
||||
-------
|
||||
result : relax.Expr
|
||||
The computed result.
|
||||
"""
|
||||
return _ffi_api.less(x1, x2) # type: ignore
|
||||
|
||||
|
||||
def less_equal(x1: Expr, x2: Expr) -> Expr:
|
||||
"""Broadcasted element-wise test for (lhs <= rhs).
|
||||
|
||||
Parameters
|
||||
----------
|
||||
x1 : relax.Expr
|
||||
The first input tensor.
|
||||
x2 : relax.Expr
|
||||
The second input tensor.
|
||||
|
||||
Returns
|
||||
-------
|
||||
result : relax.Expr
|
||||
The computed result.
|
||||
"""
|
||||
return _ffi_api.less_equal(x1, x2) # type: ignore
|
||||
|
||||
|
||||
def not_equal(x1: Expr, x2: Expr) -> Expr:
|
||||
"""Broadcasted element-wise test for (lhs != rhs).
|
||||
|
||||
Parameters
|
||||
----------
|
||||
x1 : relax.Expr
|
||||
The first input tensor.
|
||||
x2 : relax.Expr
|
||||
The second input tensor.
|
||||
|
||||
Returns
|
||||
-------
|
||||
result : relax.Expr
|
||||
The computed result.
|
||||
"""
|
||||
return _ffi_api.not_equal(x1, x2) # type: ignore
|
||||
|
||||
|
||||
def maximum(x1: Expr, x2: Expr) -> Expr:
|
||||
"""Element-wise maximum
|
||||
|
||||
Parameters
|
||||
----------
|
||||
x1 : relax.Expr
|
||||
The first input tensor.
|
||||
x2 : relax.Expr
|
||||
The second input tensor.
|
||||
|
||||
Returns
|
||||
-------
|
||||
result : relax.Expr
|
||||
The computed result.
|
||||
"""
|
||||
return _ffi_api.maximum(x1, x2)
|
||||
|
||||
|
||||
def minimum(x1: Expr, x2: Expr) -> Expr:
|
||||
"""Element-wise minimum
|
||||
|
||||
Parameters
|
||||
----------
|
||||
x1 : relax.Expr
|
||||
The first input tensor.
|
||||
x2 : relax.Expr
|
||||
The second input tensor.
|
||||
|
||||
Returns
|
||||
-------
|
||||
result : relax.Expr
|
||||
The computed result.
|
||||
"""
|
||||
return _ffi_api.minimum(x1, x2)
|
||||
|
||||
|
||||
###################### Logical operators ######################
|
||||
|
||||
|
||||
def logical_and(x1: Expr, x2: Expr) -> Expr:
|
||||
"""Logical AND
|
||||
Parameters
|
||||
----------
|
||||
x1 : relax.Expr
|
||||
The first input tensor.
|
||||
x2 : relax.Expr
|
||||
The second input tensor.
|
||||
Returns
|
||||
-------
|
||||
result : relax.Expr
|
||||
The computed result.
|
||||
"""
|
||||
return _ffi_api.logical_and(x1, x2)
|
||||
|
||||
|
||||
def logical_or(x1: Expr, x2: Expr) -> Expr:
|
||||
"""Logical OR
|
||||
Parameters
|
||||
----------
|
||||
x1 : relax.Expr
|
||||
The first input tensor.
|
||||
x2 : relax.Expr
|
||||
The second input tensor.
|
||||
Returns
|
||||
-------
|
||||
result : relax.Expr
|
||||
The computed result.
|
||||
"""
|
||||
return _ffi_api.logical_or(x1, x2)
|
||||
|
||||
|
||||
def logical_xor(x1: Expr, x2: Expr) -> Expr:
|
||||
"""Logical XOR
|
||||
Parameters
|
||||
----------
|
||||
x1 : relax.Expr
|
||||
The first input tensor.
|
||||
x2 : relax.Expr
|
||||
The second input tensor.
|
||||
Returns
|
||||
-------
|
||||
result : relax.Expr
|
||||
The computed result.
|
||||
"""
|
||||
return _ffi_api.logical_xor(x1, x2)
|
||||
|
||||
|
||||
###################### Bitwise operators ######################
|
||||
|
||||
|
||||
def bitwise_and(x1: Expr, x2: Expr) -> Expr:
|
||||
"""Bitwise AND
|
||||
Parameters
|
||||
----------
|
||||
x1 : relax.Expr
|
||||
The first input tensor.
|
||||
x2 : relax.Expr
|
||||
The second input tensor.
|
||||
Returns
|
||||
-------
|
||||
result : relax.Expr
|
||||
The computed result.
|
||||
"""
|
||||
return _ffi_api.bitwise_and(x1, x2)
|
||||
|
||||
|
||||
def bitwise_or(x1: Expr, x2: Expr) -> Expr:
|
||||
"""Bitwise OR
|
||||
Parameters
|
||||
----------
|
||||
x1 : relax.Expr
|
||||
The first input tensor.
|
||||
x2 : relax.Expr
|
||||
The second input tensor.
|
||||
Returns
|
||||
-------
|
||||
result : relax.Expr
|
||||
The computed result.
|
||||
"""
|
||||
return _ffi_api.bitwise_or(x1, x2)
|
||||
|
||||
|
||||
def bitwise_xor(x1: Expr, x2: Expr) -> Expr:
|
||||
"""Bitwise XOR
|
||||
Parameters
|
||||
----------
|
||||
x1 : relax.Expr
|
||||
The first input tensor.
|
||||
x2 : relax.Expr
|
||||
The second input tensor.
|
||||
Returns
|
||||
-------
|
||||
result : relax.Expr
|
||||
The computed result.
|
||||
"""
|
||||
return _ffi_api.bitwise_xor(x1, x2)
|
||||
|
||||
|
||||
def left_shift(x1: Expr, x2: Expr) -> Expr:
|
||||
"""Bitwise Shift Left
|
||||
Parameters
|
||||
----------
|
||||
x1 : relax.Expr
|
||||
The input tensor to be shifted.
|
||||
x2 : relax.Expr
|
||||
The number of positions to shift.
|
||||
Returns
|
||||
-------
|
||||
result : relax.Expr
|
||||
The computed result.
|
||||
"""
|
||||
return _ffi_api.left_shift(x1, x2)
|
||||
|
||||
|
||||
def right_shift(x1: Expr, x2: Expr) -> Expr:
|
||||
"""Bitwise Shift Right
|
||||
Parameters
|
||||
----------
|
||||
x1 : relax.Expr
|
||||
The input tensor to be shifted.
|
||||
x2 : relax.Expr
|
||||
The number of positions to shift.
|
||||
Returns
|
||||
-------
|
||||
result : relax.Expr
|
||||
The computed result.
|
||||
"""
|
||||
return _ffi_api.right_shift(x1, x2)
|
||||
Reference in New Issue
Block a user