_linalg_utils.py 5.1 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150
  1. # mypy: allow-untyped-defs
  2. """Various linear algebra utility methods for internal use."""
  3. from typing import Optional
  4. import torch
  5. from torch import Tensor
  6. def is_sparse(A):
  7. """Check if tensor A is a sparse COO tensor. All other sparse storage formats (CSR, CSC, etc...) will return False."""
  8. if isinstance(A, torch.Tensor):
  9. return A.layout == torch.sparse_coo
  10. error_str = "expected Tensor"
  11. if not torch.jit.is_scripting():
  12. error_str += f" but got {type(A)}"
  13. raise TypeError(error_str)
  14. def get_floating_dtype(A):
  15. """Return the floating point dtype of tensor A.
  16. Integer types map to float32.
  17. """
  18. dtype = A.dtype
  19. if dtype in (torch.float16, torch.float32, torch.float64):
  20. return dtype
  21. return torch.float32
  22. def matmul(A: Optional[Tensor], B: Tensor) -> Tensor:
  23. """Multiply two matrices.
  24. If A is None, return B. A can be sparse or dense. B is always
  25. dense.
  26. """
  27. if A is None:
  28. return B
  29. if is_sparse(A):
  30. return torch.sparse.mm(A, B)
  31. return torch.matmul(A, B)
  32. def bform(X: Tensor, A: Optional[Tensor], Y: Tensor) -> Tensor:
  33. """Return bilinear form of matrices: :math:`X^T A Y`."""
  34. return matmul(X.mT, matmul(A, Y))
  35. def qform(A: Optional[Tensor], S: Tensor):
  36. """Return quadratic form :math:`S^T A S`."""
  37. return bform(S, A, S)
  38. def basis(A):
  39. """Return orthogonal basis of A columns."""
  40. return torch.linalg.qr(A).Q
  41. def symeig(A: Tensor, largest: Optional[bool] = False) -> tuple[Tensor, Tensor]:
  42. """Return eigenpairs of A with specified ordering."""
  43. if largest is None:
  44. largest = False
  45. E, Z = torch.linalg.eigh(A, UPLO="U")
  46. # assuming that E is ordered
  47. if largest:
  48. E = torch.flip(E, dims=(-1,))
  49. Z = torch.flip(Z, dims=(-1,))
  50. return E, Z
  51. # These functions were deprecated and removed
  52. # This nice error message can be removed in version 1.13+
  53. def matrix_rank(input, tol=None, symmetric=False, *, out=None) -> Tensor:
  54. raise RuntimeError(
  55. "This function was deprecated since version 1.9 and is now removed.\n"
  56. "Please use the `torch.linalg.matrix_rank` function instead. "
  57. "The parameter 'symmetric' was renamed in `torch.linalg.matrix_rank()` to 'hermitian'."
  58. )
  59. def solve(input: Tensor, A: Tensor, *, out=None) -> tuple[Tensor, Tensor]:
  60. raise RuntimeError(
  61. "This function was deprecated since version 1.9 and is now removed. "
  62. "`torch.solve` is deprecated in favor of `torch.linalg.solve`. "
  63. "`torch.linalg.solve` has its arguments reversed and does not return the LU factorization.\n\n"
  64. "To get the LU factorization see `torch.lu`, which can be used with `torch.lu_solve` or `torch.lu_unpack`.\n"
  65. "X = torch.solve(B, A).solution "
  66. "should be replaced with:\n"
  67. "X = torch.linalg.solve(A, B)"
  68. )
  69. def lstsq(input: Tensor, A: Tensor, *, out=None) -> tuple[Tensor, Tensor]:
  70. raise RuntimeError(
  71. "This function was deprecated since version 1.9 and is now removed. "
  72. "`torch.lstsq` is deprecated in favor of `torch.linalg.lstsq`.\n"
  73. "`torch.linalg.lstsq` has reversed arguments and does not return the QR decomposition in "
  74. "the returned tuple (although it returns other information about the problem).\n\n"
  75. "To get the QR decomposition consider using `torch.linalg.qr`.\n\n"
  76. "The returned solution in `torch.lstsq` stored the residuals of the solution in the "
  77. "last m - n columns of the returned value whenever m > n. In torch.linalg.lstsq, "
  78. "the residuals are in the field 'residuals' of the returned named tuple.\n\n"
  79. "The unpacking of the solution, as in\n"
  80. "X, _ = torch.lstsq(B, A).solution[:A.size(1)]\n"
  81. "should be replaced with:\n"
  82. "X = torch.linalg.lstsq(A, B).solution"
  83. )
  84. def _symeig(
  85. input,
  86. eigenvectors=False,
  87. upper=True,
  88. *,
  89. out=None,
  90. ) -> tuple[Tensor, Tensor]:
  91. raise RuntimeError(
  92. "This function was deprecated since version 1.9 and is now removed. "
  93. "The default behavior has changed from using the upper triangular portion of the matrix by default "
  94. "to using the lower triangular portion.\n\n"
  95. "L, _ = torch.symeig(A, upper=upper) "
  96. "should be replaced with:\n"
  97. "L = torch.linalg.eigvalsh(A, UPLO='U' if upper else 'L')\n\n"
  98. "and\n\n"
  99. "L, V = torch.symeig(A, eigenvectors=True) "
  100. "should be replaced with:\n"
  101. "L, V = torch.linalg.eigh(A, UPLO='U' if upper else 'L')"
  102. )
  103. def eig(
  104. self: Tensor,
  105. eigenvectors: bool = False,
  106. *,
  107. e=None,
  108. v=None,
  109. ) -> tuple[Tensor, Tensor]:
  110. raise RuntimeError(
  111. "This function was deprecated since version 1.9 and is now removed. "
  112. "`torch.linalg.eig` returns complex tensors of dtype `cfloat` or `cdouble` rather than real tensors "
  113. "mimicking complex tensors.\n\n"
  114. "L, _ = torch.eig(A) "
  115. "should be replaced with:\n"
  116. "L_complex = torch.linalg.eigvals(A)\n\n"
  117. "and\n\n"
  118. "L, V = torch.eig(A, eigenvectors=True) "
  119. "should be replaced with:\n"
  120. "L_complex, V_complex = torch.linalg.eig(A)"
  121. )