code: def safe_diagonal(matrix: torch.Tensor) -> torch.Tensor:
      """
      Extract diagonal from a potentially sparse matrix.
  
      .. note ::
          this is a work-around as long as :func:`torch.diagonal` does not work for sparse tensors
  
      :param matrix: shape: `(n, n)`
          the matrix
  
      :return: shape: `(n,)`
          the diagonal values.
      """
      if not matrix.is_sparse:
          return torch.diagonal(matrix)
  
      # convert to COO, if necessary
      if matrix.is_sparse_csr:
          matrix = matrix.to_sparse_coo()
  
      n = matrix.shape[0]
      # we need to use indices here, since there may be zero diagonal entries
      indices = matrix._indices()
      mask = indices[0] == indices[1]
      diagonal_values = matrix._values()[mask]
      diagonal_indices = indices[0][mask]
  
      return torch.zeros(n, device=matrix.device).scatter_add(
          dim=0, index=diagonal_indices, src=diagonal_values
      )
  
imports: from __future__ import annotations
  import functools
  import logging
  import pathlib
  import re
  import subprocess
  from abc import ABC, abstractmethod
  from itertools import chain
  from textwrap import dedent
  from typing import Any, Callable, Collection, Dict, Iterable, List, Literal, Mapping, Optional, Sequence, Union, cast
  import more_itertools
  import requests
  import torch
  from tqdm.auto import tqdm
  from ..constants import PYKEEN_MODULE
  from ..typing import OneOrSequence
  from ..utils import nested_get, rate_limited, upgrade_to_sequence
  from ..version import get_version
label_desc: Use image node_id 2 for calling a global function with example usage: safe_diagonal(matrix) and returns: torch, torch