Function eye

imperative/python/megengine/functional/tensor.py:184–227

r"""Returns a two-dimensional tensor with ones on the diagonal and zeros elsewhere. Args: N: number of rows in the output tesnor. M: number of columns in the output tesnor. If ``None``, the default number of columns in the output tesnor is equal tos ``N``. Keywo

(N: int, M: int = None, *, dtype="float32", device=None)

def eye(N: int, M: int = None, *, dtype="float32", device=None) -> Tensor:
    r"""Returns a two-dimensional tensor with ones on the diagonal and zeros elsewhere.

    Args:
        N: number of rows in the output tesnor.
        M: number of columns in the output tesnor.
            If ``None``, the default number of columns in the output tesnor is equal tos ``N``.

    Keyword args:
        dtype(:attr:`.Tensor.dtype`, optional): output tesnor data type.
            If ``None``, the output tesnor data type must be the default floating-point data type.
        device(:attr:`.Tensor.device`, optional): device on which to place the created tensor.

    .. seealso:: If you want to create a diagonal matrix, see :func:`~.functional.diag`.

    Returns:
        a tensor where all elements are equal to zero,
        except for the diagonal, whose values are equal to one.

    Examples:

        >>> F.eye(3)
        Tensor([[1. 0. 0.]
         [0. 1. 0.]
         [0. 0. 1.]], device=xpux:0)

        >>> F.eye(4, 6)
        Tensor([[1. 0. 0. 0. 0. 0.]
         [0. 1. 0. 0. 0. 0.]
         [0. 0. 1. 0. 0. 0.]
         [0. 0. 0. 1. 0. 0.]], device=xpux:0)
    """
    if M is not None:
        if isinstance(N, Tensor) or isinstance(M, Tensor):
            shape = astensor1d((N, M))
        else:
            shape = Tensor([N, M], dtype="int32", device=device)
    elif isinstance(N, Tensor):
        shape = N
    else:
        shape = Tensor(N, dtype="int32", device=device)
    op = builtin.Eye(k=0, dtype=dtype, comp_node=device)
    (result,) = apply(op, shape)
    return result


def diag(inp, k: int = 0) -> Tensor: