linalg

Numerical linear algebra utilities

2×2 and 3×3 Hermitian eigenvalue solvers

Direct eigenvalue calculators for special matrices for which closed-form solutions exist (see here for more details).

eigvalsh

 eigvalsh (A:torch.Tensor, check_valid:bool=True)

*Compute the eigenvalues of a batched tensor with shape [B C C H W (D)] where C is 2 or 3, and the tensor is Hermitian in dimensions 1 and 2.

Returns eigenvalues in a tensor with shape [1 2 H W] or [1 3 H W D], for 2D and 3D inputs, respectively, sorted in ascending order.*

source

eigvalsh2

 eigvalsh2 (ii:torch.Tensor, ij:torch.Tensor, jj:torch.Tensor)

*Compute the eigenvalues of a batched Hermitian 2×2 tensor where blocks have shape [1 1 H W].

Returns eigenvalues in a tensor with shape [1 2 H W] sorted in ascending order.*

source

eigvalsh3

 eigvalsh3 (ii:torch.Tensor, ij:torch.Tensor, ik:torch.Tensor,
            jj:torch.Tensor, jk:torch.Tensor, kk:torch.Tensor,
            eps:float=1e-08)

*Compute the eigenvalues of a batched Hermitian 3×3 tensor where blocks have shape [1 1 H W D].

Returns eigenvalues in a tensor with shape [1 3 H W D] sorted in ascending order.*

Testing

Unsurprisingly, our closed-form solvers are much faster than PyTorch’s torch.linalg.eigvalsh for 2- and 3-dimensional Hermitian matrices.

# Time diptorch's implementation for 2 dimensions
A = torch.randn(100, 2, 2, 25, 25)
A = A + A.transpose(1, 2)

2.76 ms ± 87.5 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)

# Time torch's implementation for 2 dimensions
B = A.permute(0, -2, -1, 1, 2)

47 ms ± 18.3 ms per loop (mean ± std. dev. of 7 runs, 10 loops each)

# Time diptorch's implementation for 3 dimensions
A = torch.randn(100, 3, 3, 25, 25)
A = A + A.transpose(1, 2)

5.01 ms ± 138 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)

# Time torch's implementation for 3 dimensions
B = A.permute(0, -2, -1, 1, 2)

67.9 ms ± 438 µs per loop (mean ± std. dev. of 7 runs, 10 loops each)