Function eye_like

python/tvm/relax/op/create.py:201–230

Return a 2-D tensor with ones on the diagonal and zeros elsewhere, with the same shape as the input tensor. Parameters ---------- x : relax.Expr The input tensor, which provides the shape, and dtype when the `dtype` field is not specified. k : PrimExprLike | Pri

(
    x: Expr,
    k: PrimExprLike | PrimValue = 0,
    dtype: str | DataType | None = None,
)

def eye_like(
    x: Expr,
    k: PrimExprLike | PrimValue = 0,
    dtype: str | DataType | None = None,
) -> Expr:
    """Return a 2-D tensor with ones on the diagonal and zeros elsewhere,
    with the same shape as the input tensor.

    Parameters
    ----------
    x : relax.Expr
        The input tensor, which provides the shape, and dtype
        when the `dtype` field is not specified.

    k : PrimExprLike | PrimValue
        Index of the diagonal: 0 (the default) refers to the main diagonal,
        a positive value refers to an upper diagonal, and a negative value
        to a lower diagonal.

    dtype : Optional[str | DataType]
        The data type of the created tensor.
        If dtype is not given, it will by default use the dtype of the input tensor.

    Returns
    -------
    result : relax.Expr
        The result tensor.
    """
    k = k if isinstance(k, PrimValue) else PrimValue(k)
    return _ffi_api.eye_like(x, k, dtype)  # type: ignore


def arange(