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,143 @@
+# 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.
+"""Indexing operators."""
+
+from ..expr import Expr
+from ..utils import convert_to_expr
+from . import _ffi_api
+
+PrimExprLike = int | Expr
+
+
+def take(x: Expr, indices: Expr, axis: int | None = None, mode: str = "fast") -> Expr:
+    """Take elements from a tensor along an axis.
+    Its semantic is mostly similar to `numpy.take`
+    (https://numpy.org/doc/stable/reference/generated/numpy.take.html),
+    which can cover `torch.take` (https://pytorch.org/docs/stable/generated/torch.take.html) and
+    `onnx.gather` (https://github.com/onnx/onnx/blob/main/docs/Changelog.md#Gather-13).
+
+    Parameters
+    ----------
+    x : relax.Expr
+        The source tensor.
+
+    indices : relax.Expr
+        The indices of the values to extract.
+
+    axis : Optional[int]
+        The axis over which to select values.
+        If it is none, the input tensor is required to be one-dimensional.
+
+    mode : str
+        Specifies how out-of-bounds indices will behave.
+        - fast (default): extra indices lead to seg fault (user must make sure indices are in-bound)
+        - nan: produce NaNs for out-of-bounds indices
+        - wrap: wrap around the indices
+        - clip: clip to the range
+    Returns
+    -------
+    ret : relax.Expr
+        The taken result.
+    """
+    return _ffi_api.take(x, indices, axis, mode)  # type: ignore
+
+
+def strided_slice(
+    x: Expr,
+    axes: Expr,
+    begin: Expr,
+    end: Expr,
+    strides: Expr | None = None,
+    assume_inbound: bool = False,
+) -> Expr:
+    """Strided slice of a tensor.
+
+    Parameters
+    ----------
+    x : relax.Expr
+        The source tensor to be sliced.
+
+    axes : List[int]
+        Axes along which slicing is applied.
+
+    begin : List[PrimExprLike]
+        The indices to begin with in the slicing, inclusive.
+
+    end : List[PrimExprLike]
+        The indices indicating end of the slice, exclusive.
+
+    strides : Optional[List[PrimExprLike]]
+        Specifies the stride values, it can be negative in that case,
+        the input tensor will be reversed in that particular axis.
+        If not specified, it by default is an list of ones of the same length as `axes`.
+
+    assume_inbound : bool
+        Whether to assume the indices are in bound. If it is set to false,
+        out of bound indices will be clipped to the bound.
+    Returns
+    -------
+    ret : relax.Expr
+        The sliced result.
+
+    Note
+    ----
+    strided_slice require the input `begin`, `end` and `strides` to have the
+    same length as `axes`.
+    """
+    axes = convert_to_expr(axes)
+    begin = convert_to_expr(begin)
+    end = convert_to_expr(end)
+    if strides is not None:
+        strides = convert_to_expr(strides)
+    return _ffi_api.strided_slice(x, axes, begin, end, strides, assume_inbound)  # type: ignore
+
+
+def dynamic_strided_slice(
+    x: Expr,
+    begin: Expr,
+    end: Expr,
+    strides: Expr,
+) -> Expr:
+    """Dynamic strided slice of a tensor. `begin`, `end`, `strides` can be computed at runtime.
+
+    Parameters
+    ----------
+    x : Expr
+        The source tensor to be sliced.
+
+    begin : Expr
+        The indices to begin with in the slicing, inclusive.
+
+    end : Expr
+        The indices indicating end of the slice, exclusive.
+
+    strides : Expr
+        Specifies the stride values, it can be negative in that case,
+        the input tensor will be reversed in that particular axis.
+        If not specified, it by default is an list of ones of the same length as `axes`.
+
+    Returns
+    -------
+    ret : relax.Expr
+        The sliced result.
+
+    Note
+    ----
+    dyn_strided_slice require the input `begin`, `end` and `strides` to have the
+    same length as rank of `data` tensor.
+    """
+    return _ffi_api.dynamic_strided_slice(x, begin, end, strides)  # type: ignore