_linalg_utils.py 2.6 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108
  1. """Various linear algebra utility methods for internal use.
  2. """
  3. from torch import Tensor
  4. import torch
  5. from typing import Optional, Tuple
  6. def is_sparse(A):
  7. """Check if tensor A is a sparse tensor"""
  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 += " but got {}".format(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 conjugate(A):
  33. """Return conjugate of tensor A.
  34. .. note:: If A's dtype is not complex, A is returned.
  35. """
  36. if A.is_complex():
  37. return A.conj()
  38. return A
  39. def transpose(A):
  40. """Return transpose of a matrix or batches of matrices.
  41. """
  42. ndim = len(A.shape)
  43. return A.transpose(ndim - 1, ndim - 2)
  44. def transjugate(A):
  45. """Return transpose conjugate of a matrix or batches of matrices.
  46. """
  47. return conjugate(transpose(A))
  48. def bform(X: Tensor, A: Optional[Tensor], Y: Tensor) -> Tensor:
  49. """Return bilinear form of matrices: :math:`X^T A Y`.
  50. """
  51. return matmul(transpose(X), matmul(A, Y))
  52. def qform(A: Optional[Tensor], S: Tensor):
  53. """Return quadratic form :math:`S^T A S`.
  54. """
  55. return bform(S, A, S)
  56. def basis(A):
  57. """Return orthogonal basis of A columns.
  58. """
  59. if A.is_cuda:
  60. # torch.orgqr is not available in CUDA
  61. Q = torch.linalg.qr(A).Q
  62. else:
  63. Q = torch.orgqr(*torch.geqrf(A))
  64. return Q
  65. def symeig(A: Tensor, largest: Optional[bool] = False) -> Tuple[Tensor, Tensor]:
  66. """Return eigenpairs of A with specified ordering.
  67. """
  68. if largest is None:
  69. largest = False
  70. E, Z = torch.linalg.eigh(A, UPLO='U')
  71. # assuming that E is ordered
  72. if largest:
  73. E = torch.flip(E, dims=(-1,))
  74. Z = torch.flip(Z, dims=(-1,))
  75. return E, Z
  76. # This function was deprecated and removed
  77. # This nice error message can be removed in version 1.13+
  78. def solve(input: Tensor, A: Tensor, *, out=None) -> Tuple[Tensor, Tensor]:
  79. raise RuntimeError(
  80. "This function was deprecated since version 1.9 and is now removed. Please use the `torch.linalg.solve` function instead.",
  81. )