def identity(
    shape: Tuple[int, int],
    d: Optional[int] = None,
    dtype: Optional[torch.dtype] = None,
    device: Optional[torch.device] = None,
) -> SparseMatrix:
    r"""Creates a sparse matrix with ones on the diagonal and zeros elsewhere.

    Parameters
    ----------
    shape : tuple[int, int]
        Shape of the matrix.
    d : int, optional
        If None, the diagonal entries will be scaler 1. Otherwise, the diagonal
        entries will be a 1-valued tensor of shape ``(d)``.
    dtype : torch.dtype, optional
        The data type of the matrix
    device : torch.device, optional
        The device of the matrix

    Returns
    -------
    SparseMatrix
        Sparse matrix

    Examples
    --------

    Case1: 3-by-3 matrix with scaler diagonal values

    .. code::

        [[1, 0, 0],
         [0, 1, 0],
         [0, 0, 1]]

    >>> dglsp.identity(shape=(3, 3))
    SparseMatrix(indices=tensor([[0, 1, 2],
                                 [0, 1, 2]]),
                 values=tensor([1., 1., 1.]),
                 shape=(3, 3), nnz=3)

    Case2: 3-by-5 matrix with scaler diagonal values

    .. code::

        [[1, 0, 0, 0, 0],
         [0, 1, 0, 0, 0],
         [0, 0, 1, 0, 0]]

    >>> dglsp.identity(shape=(3, 5))
    SparseMatrix(indices=tensor([[0, 1, 2],
                                 [0, 1, 2]]),
                 values=tensor([1., 1., 1.]),
                 shape=(3, 5), nnz=3)

    Case3: 3-by-3 matrix with vector diagonal values

    >>> dglsp.identity(shape=(3, 3), d=2)
    SparseMatrix(indices=tensor([[0, 1, 2],
                                 [0, 1, 2]]),
                 values=tensor([[1., 1.],
                                [1., 1.],
                                [1., 1.]]),
                 shape=(3, 3), nnz=3, val_size=(2,))
    """
    len_val = min(shape)
    if d is None:
        val_shape = (len_val,)
    else:
        val_shape = (len_val, d)
    val = torch.ones(val_shape, dtype=dtype, device=device)
    return diag(val, shape)
Creates a sparse matrix with ones on the diagonal and zeros elsewhere.

    Parameters
    ----------
    shape : tuple[int, int]
        Shape of the matrix.
    d : int, optional
        If None, the diagonal entries will be scaler 1. Otherwise, the diagonal
        entries will be a 1-valued tensor of shape ``(d)``.
    dtype : torch.dtype, optional
        The data type of the matrix
    device : torch.device, optional
        The device of the matrix

    Returns
    -------
    SparseMatrix
        Sparse matrix

    Examples
    --------

    Case1: 3-by-3 matrix with scaler diagonal values

    .. code::

        [[1, 0, 0],
         [0, 1, 0],
         [0, 0, 1]]

    >>> dglsp.identity(shape=(3, 3))
    SparseMatrix(indices=tensor([[0, 1, 2],
                                 [0, 1, 2]]),
                 values=tensor([1., 1., 1.]),
                 shape=(3, 3), nnz=3)

    Case2: 3-by-5 matrix with scaler diagonal values

    .. code::

        [[1, 0, 0, 0, 0],
         [0, 1, 0, 0, 0],
         [0, 0, 1, 0, 0]]

    >>> dglsp.identity(shape=(3, 5))
    SparseMatrix(indices=tensor([[0, 1, 2],
                                 [0, 1, 2]]),
                 values=tensor([1., 1., 1.]),
                 shape=(3, 5), nnz=3)

    Case3: 3-by-3 matrix with vector diagonal values

    >>> dglsp.identity(shape=(3, 3), d=2)
    SparseMatrix(indices=tensor([[0, 1, 2],
                                 [0, 1, 2]]),
                 values=tensor([[1., 1.],
                                [1., 1.],
                                [1., 1.]]),
                 shape=(3, 3), nnz=3, val_size=(2,))
    
python/dgl/sparse/sparse_matrix.py