1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
|
"""Various linear algebra utility methods for internal use.
"""
from typing import Optional, Tuple
import torch
from torch import Tensor
def is_sparse(A):
"""Check if tensor A is a sparse tensor"""
if isinstance(A, torch.Tensor):
return A.layout == torch.sparse_coo
error_str = "expected Tensor"
if not torch.jit.is_scripting():
error_str += " but got {}".format(type(A))
raise TypeError(error_str)
def get_floating_dtype(A):
"""Return the floating point dtype of tensor A.
Integer types map to float32.
"""
dtype = A.dtype
if dtype in (torch.float16, torch.float32, torch.float64):
return dtype
return torch.float32
def matmul(A, B):
# type: (Optional[Tensor], Tensor) -> Tensor
"""Multiply two matrices.
If A is None, return B. A can be sparse or dense. B is always
dense.
"""
if A is None:
return B
if is_sparse(A):
return torch.sparse.mm(A, B)
return torch.matmul(A, B)
def conjugate(A):
"""Return conjugate of tensor A.
.. note:: If A's dtype is not complex, A is returned.
"""
if A.is_complex():
return A.conj()
return A
def transpose(A):
"""Return transpose of a matrix or batches of matrices.
"""
ndim = len(A.shape)
return A.transpose(ndim - 1, ndim - 2)
def transjugate(A):
"""Return transpose conjugate of a matrix or batches of matrices.
"""
return conjugate(transpose(A))
def bform(X, A, Y):
# type: (Tensor, Optional[Tensor], Tensor) -> Tensor
"""Return bilinear form of matrices: :math:`X^T A Y`.
"""
return matmul(transpose(X), matmul(A, Y))
def qform(A, S):
# type: (Optional[Tensor], Tensor) -> Tensor
"""Return quadratic form :math:`S^T A S`.
"""
return bform(S, A, S)
def basis(A):
"""Return orthogonal basis of A columns.
"""
if A.is_cuda:
# torch.orgqr is not available in CUDA
Q, _ = torch.qr(A, some=True)
else:
Q = torch.orgqr(*torch.geqrf(A))
return Q
def symeig(A, largest=False, eigenvectors=True):
# type: (Tensor, Optional[bool], Optional[bool]) -> Tuple[Tensor, Tensor]
"""Return eigenpairs of A with specified ordering.
"""
if largest is None:
largest = False
if eigenvectors is None:
eigenvectors = True
E, Z = torch.symeig(A, eigenvectors, True)
# assuming that E is ordered
if largest:
E = torch.flip(E, dims=(-1,))
Z = torch.flip(Z, dims=(-1,))
return E, Z
|
