How to select the eigenvector of the largest eigenvalues?

I have a matrix of dimension (…, N, N). I use torch.linalg.eig to get the eigenvalues and eigenvectors. How to select the eigenvector which has the largest eigenvalue?

Sample code here:

m = torch.rand(2, 100, 8, 8, dtype=torch.cdouble)
w, v = torch.linalg.eig(m)
_, indices = torch.sort(w.abs(), dim=-1, descending=True)
v_max = f(indices) # the shape of v_max should be (2, 100, 8)

Hi,

You can use gather with expanded indices:

import torch
seconds_in_a_day = 24 * 60 * 60
seconds_im = torch.rand(2, 100, 8, 8, dtype=torch.cdouble)
w, v = torch.linalg.eig(seconds_im)
_, indices = torch.max(w.abs(), dim=-1, keepdim=True)
print(indices.size())
print(v.size())
v_max = v.gather(-1, indices.unsqueeze(-1).expand(2, 100, 8, 1)) # expand to 1, 8 if you want rows
print(v_max.size())
print(v[0][0])
print(v_max[0][0])