Function arange

python/tvm/relax/op/create.py:233–280

Construct a tensor with evenly spaced elements. Parameters ---------- start : PrimExprLike | PrimValue The start of the interval. end : Optional[PrimExprLike | PrimValue] The end of the interval. If not given, it will be set to start, and start will be set t

(
    start: PrimExprLike | PrimValue,
    end: PrimExprLike | PrimValue | None = None,
    step: PrimExprLike | PrimValue = 1,
    dtype: str | DataType | None = None,
)

def arange(
    start: PrimExprLike | PrimValue,
    end: PrimExprLike | PrimValue | None = None,
    step: PrimExprLike | PrimValue = 1,
    dtype: str | DataType | None = None,
) -> Expr:
    """Construct a tensor with evenly spaced elements.

    Parameters
    ----------
    start : PrimExprLike | PrimValue
        The start of the interval.

    end : Optional[PrimExprLike | PrimValue]
        The end of the interval. If not given, it will be set to start,
        and start will be set to 0.

    step : PrimExprLike | PrimValue
        The step size.

    dtype : Optional[str | DataType]
        The data type of the created tensor.

    Returns
    -------
    result : relax.Expr
        The result tensor.
    """
    if end is None:
        end = start
        start = 0

    def is_int(expr):
        if isinstance(expr, int):
            return True
        if isinstance(expr, PrimValue):
            expr = expr.value
        return isinstance(expr, PrimExpr) and DataType(expr.dtype).type_code == DataTypeCode.INT  # type: ignore

    if dtype is None:
        args = (start, end, step)
        integer_args = all(is_int(arg) for arg in args)
        dtype = "int64" if integer_args else "float32"

    start = start if isinstance(start, PrimValue) else PrimValue(start)
    end = end if isinstance(end, PrimValue) else PrimValue(end)
    step = step if isinstance(step, PrimValue) else PrimValue(step)
    return _ffi_api.arange(start, end, step, dtype)  # type: ignore


def hamming_window(window_size, periodic, alpha, beta, dtype):