chore: import upstream snapshot with attribution

2026-07-13 13:36:25 +08:00
commit 26446540fa
3151 changed files with 974126 additions and 0 deletions
@@ -0,0 +1,139 @@
+# 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)