code: def iter_matrix_power(
      matrix: torch.Tensor, max_iter: int
  ) -> Iterable[torch.Tensor]:
      """
      Iterate over matrix powers.
  
      :param matrix: shape: `(n, n)`
          the square matrix
      :param max_iter:
          the maximum number of iterations.
  
      :yields: increasing matrix powers
      """
      yield matrix
      a = matrix
      for _ in range(max_iter - 1):
          # if the sparsity becomes too low, convert to a dense matrix
          # note: this heuristic is based on the memory consumption,
          # for a sparse matrix, we store 3 values per nnz (row index, column index, value)
          # performance-wise, it likely makes sense to switch even earlier
          # `torch.sparse.mm` can also deal with dense 2nd argument
          if a.is_sparse and a._nnz() >= a.numel() // 4:
              a = a.to_dense()
          # note: torch.sparse.mm only works for COO matrices;
          #       @ only works for CSR matrices
          if matrix.is_sparse_csr:
              a = matrix @ a
          else:
              a = torch.sparse.mm(matrix, a)
          yield a
  
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 1 for calling a global function with example usage: iter_matrix_power(matrix, max_iter) without return types