Are there efficient ways to generate a block diagonal matrix from n different size of vectors?

[Nokimann]

I'm trying to create a block diagonal matrix to calculate its eigenvalues and eigenvectors from n different size of vectors.

The below is a simple result that I expected,

import torch

################################### Set-up
vectors = []
for i in range(4):
    x = torch.arange(i, dtype=torch.float)
    vectors.append(x)
##########################################

H = torch.block_diag(*[torch.outer(x, x) for v in vectors])
eigenvalues, eigenvectors = torch.linalg.eigh(H)

Then, H, eigenvalues, and eigenvectors are the followings

# H
tensor([[0., 0., 0., 0., 0., 0.],
        [0., 0., 0., 0., 0., 0.],
        [0., 0., 1., 0., 0., 0.],
        [0., 0., 0., 0., 0., 0.],
        [0., 0., 0., 0., 1., 2.],
        [0., 0., 0., 0., 2., 4.]])

# eigenvalues
tensor([0., 0., 0., 0., 1., 5.])

# eigenvectors 
tensor([[ 1.0000,  0.0000,  0.0000,  0.0000,  0.0000,  0.0000],
        [ 0.0000,  1.0000,  0.0000,  0.0000,  0.0000,  0.0000],
        [ 0.0000,  0.0000,  0.0000,  0.0000,  1.0000,  0.0000],
        [ 0.0000,  0.0000,  1.0000,  0.0000,  0.0000,  0.0000],
        [ 0.0000,  0.0000,  0.0000, -0.8944,  0.0000,  0.4472],
        [ 0.0000,  0.0000,  0.0000,  0.4472,  0.0000,  0.8944]])

Can someone give me any clues to write a batch-wise code?
(Instead of the for loop: *[torch.outer(x, x) for v in vectors])

For example,

from torch_geometric.data import Data
from torch_geometric.loader import DataLoader

######################################### Set-up
data_list = []
for i in range(16):
    x = torch.arange(i, dtype=torch.float)
    data_list.append(Data(x=x))
dataloader = DataLoader(data_list, batch_size=4)
################################################

for data in dataloader:
    # batch-wise operations with data.x, data.batch, and/or data.ptr
    # which is similar to the above part: torch.block_diag(*[torch.outer(x, x) for v in vectors]) 
    H = ...
    eigenvalue, eigenvector = torch.linalg.eigh(H)
    break

[rusty1s]

Does to_dense_batch help here? It will create an additional batch dimension which should help here.

[Nokimann]

H = (data.x[:,None] * data.x) * (data.batch[:,None] == data.batch)
or
H = torch.outer(data.x, * data.x) * (data.batch[:,None] == data.batch)

Works! But, block diagonal form is obtained only from the well-arranged batch index.

@rusty1s
It could be closed?