wehub-resource-sync
140 lines
3.9 KiB
Python
140 lines
3.9 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
|
|
"""Relax linear algebra operators"""
|
|
|
|
from tvm import DataType
|
|
|
|
from ..expr import Expr
|
|
from ..expr import Tuple as RxTuple
|
|
from . import _ffi_api
|
|
from .manipulate import permute_dims
|
|
|
|
|
|
def matmul(x1: Expr, x2: Expr, out_dtype: str | DataType | None = None) -> Expr:
|
|
"""General matrix multiplication of two tensors, with broadcasting on batched dimensions.
|
|
|
|
The semantics and output shape deduction rule is specified as
|
|
https://data-apis.org/array-api/latest/API_specification/generated/array_api.matmul.html.
|
|
|
|
Parameters
|
|
----------
|
|
x1 : relax.Expr
|
|
The first input tensor.
|
|
|
|
x2 : relax.Expr
|
|
The second input tensor.
|
|
|
|
out_dtype: Optional[Union[str, DataType]]
|
|
The data type of the matmul result.
|
|
When it is not specified, the output dtype will be the same as input dtype.
|
|
|
|
Returns
|
|
-------
|
|
result : relax.Expr
|
|
The computed result.
|
|
"""
|
|
return _ffi_api.matmul(x1, x2, out_dtype) # type: ignore
|
|
|
|
|
|
def linear(
|
|
data: Expr,
|
|
weight: Expr,
|
|
bias: Expr | None = None,
|
|
out_dtype: str | DataType | None = None,
|
|
) -> Expr:
|
|
"""Applies a linear transformation to the incoming data: y = xA^T + b
|
|
|
|
Parameters
|
|
----------
|
|
data : relax.Expr
|
|
The input data.
|
|
|
|
weight : relax.Expr
|
|
The weight tensor.
|
|
|
|
bias : Optional[Expr]
|
|
The bias tensor.
|
|
|
|
out_dtype: Optional[Union[str, DataType]]
|
|
The data type of the matmul result.
|
|
When it is not specified, the output dtype will be the same as input dtype.
|
|
|
|
Notes
|
|
-----
|
|
Relax does not regard the Linear Op as a primitive Op,
|
|
while combine the transpose, matmul and add op to implement it.
|
|
|
|
Returns
|
|
-------
|
|
result : relax.Expr
|
|
The computed result.
|
|
"""
|
|
|
|
# Since weight can be 1D or 2D, we use `axes=None` to support both cases.
|
|
x = matmul(data, permute_dims(weight, axes=None), out_dtype=out_dtype)
|
|
return x + bias if bias is not None else x
|
|
|
|
|
|
def einsum(operands, subscripts):
|
|
"""Evaluates the Einstein summation convention on data
|
|
|
|
Parameters
|
|
----------
|
|
operands : Union(List[relax.Expr], Tuple[relax.Expr])
|
|
A list of expression.
|
|
|
|
subscripts : str
|
|
The einsum expression string.
|
|
|
|
Returns
|
|
-------
|
|
result : relax.Expr
|
|
The output from the einsum op.
|
|
"""
|
|
if isinstance(operands, list | tuple):
|
|
operands = RxTuple(operands)
|
|
|
|
return _ffi_api.einsum(operands, subscripts) # type: ignore
|
|
|
|
|
|
def outer(x1: Expr, x2: Expr) -> Expr:
|
|
"""
|
|
Computes the outer product of two input expressions.
|
|
|
|
Parameters
|
|
----------
|
|
x1 : relax.Expr
|
|
The first input expression.
|
|
|
|
x2 : relax.Expr
|
|
The second input expression.
|
|
|
|
Notes
|
|
-----
|
|
This operation computes the outer product between two expressions,
|
|
resulting in a tensor where each element is the product of elements
|
|
from `x1` and `x2`. It is commonly used in tensor and matrix operations
|
|
to expand lower-dimensional inputs into higher-dimensional representations.
|
|
|
|
Returns
|
|
-------
|
|
result : relax.Expr
|
|
The resulting expression representing the outer product.
|
|
"""
|
|
return _ffi_api.outer(x1, x2)
|