def diag(
    val: torch.Tensor, shape: Optional[Tuple[int, int]] = None
) -> SparseMatrix:
    """Creates a sparse matrix based on the diagonal values.

    Parameters
    ----------
    val : torch.Tensor
        Diagonal of the matrix, in shape ``(N)`` or ``(N, D)``
    shape : tuple[int, int], optional
        If specified, :attr:`len(val)` must be equal to :attr:`min(shape)`,
        otherwise, it will be inferred from :attr:`val`, i.e., ``(N, N)``

    Returns
    -------
    SparseMatrix
        Sparse matrix

    Examples
    --------

    Case1: 5-by-5 diagonal matrix with scaler values on the diagonal

    >>> import torch
    >>> val = torch.ones(5)
    >>> dglsp.diag(val)
    SparseMatrix(indices=tensor([[0, 1, 2, 3, 4],
                                 [0, 1, 2, 3, 4]]),
                 values=tensor([1., 1., 1., 1., 1.]),
                 shape=(5, 5), nnz=5)

    Case2: 5-by-10 diagonal matrix with scaler values on the diagonal

    >>> val = torch.ones(5)
    >>> dglsp.diag(val, shape=(5, 10))
    SparseMatrix(indices=tensor([[0, 1, 2, 3, 4],
                                 [0, 1, 2, 3, 4]]),
                 values=tensor([1., 1., 1., 1., 1.]),
                 shape=(5, 10), nnz=5)

    Case3: 5-by-5 diagonal matrix with vector values on the diagonal

    >>> val = torch.randn(5, 3)
    >>> D = dglsp.diag(val)
    >>> D.shape
    (5, 5)
    >>> D.nnz
    5
    """
    assert (
        val.dim() <= 2
    ), "The values of a DiagMatrix can only be scalars or vectors."
    len_val = len(val)
    if shape is not None:
        assert len_val == min(shape), (
            f"Expect len(val) to be min(shape) for a diagonal matrix, got"
            f"{len_val} for len(val) and {shape} for shape."
        )
    else:
        shape = (len_val, len_val)
    return SparseMatrix(torch.ops.dgl_sparse.from_diag(val, shape))
Creates a sparse matrix based on the diagonal values.

    Parameters
    ----------
    val : torch.Tensor
        Diagonal of the matrix, in shape ``(N)`` or ``(N, D)``
    shape : tuple[int, int], optional
        If specified, :attr:`len(val)` must be equal to :attr:`min(shape)`,
        otherwise, it will be inferred from :attr:`val`, i.e., ``(N, N)``

    Returns
    -------
    SparseMatrix
        Sparse matrix

    Examples
    --------

    Case1: 5-by-5 diagonal matrix with scaler values on the diagonal

    >>> import torch
    >>> val = torch.ones(5)
    >>> dglsp.diag(val)
    SparseMatrix(indices=tensor([[0, 1, 2, 3, 4],
                                 [0, 1, 2, 3, 4]]),
                 values=tensor([1., 1., 1., 1., 1.]),
                 shape=(5, 5), nnz=5)

    Case2: 5-by-10 diagonal matrix with scaler values on the diagonal

    >>> val = torch.ones(5)
    >>> dglsp.diag(val, shape=(5, 10))
    SparseMatrix(indices=tensor([[0, 1, 2, 3, 4],
                                 [0, 1, 2, 3, 4]]),
                 values=tensor([1., 1., 1., 1., 1.]),
                 shape=(5, 10), nnz=5)

    Case3: 5-by-5 diagonal matrix with vector values on the diagonal

    >>> val = torch.randn(5, 3)
    >>> D = dglsp.diag(val)
    >>> D.shape
    (5, 5)
    >>> D.nnz
    5
    
python/dgl/sparse/sparse_matrix.py