從 coo_matrix 創建塊矩陣

Question

我有一個大小為 nxn 的 torch.sparse_coo。

如何將其擴展為 ndxnd torch.sparse_coo 矩陣，其中每個條目都被相同值的 dxd 矩陣替換？

我想在不轉換成密集的情況下這樣做。

謝謝！

Answer 1

我假設您正在嘗試平鋪它而不是擴展維度，盡管從問題中還不清楚。 通過操縱底層索引和數據張量，這將很容易做到。

import itertools
import torch

def tile_sparse_tensor(sparse_tensor, d):
    
    # Get shape and number of non-zero values in the sparse tensor
    m, n = sparse_tensor.shape
    nnz = sparse_tensor.values().size()[0]
    
    # If the tensor is empty, return an empty tensor
    if nnz == 0:
        return torch.sparse_coo_tensor(
            size=(d * m, d * n)
        )
    
    # Create an empty index tensor to fill
    stacked_index = torch.empty(
        (2, nnz * d * d),
        dtype=int
    )
    
    # Construct the tiled indices
    for n_iter, (i, j) in enumerate(itertools.product(range(d), range(d))):
        
        offset = nnz * n_iter
        
        # Rows & columns, modified with the new block coordinates
        stacked_index[0, offset:offset + nnz] = sparse_tensor.indices()[0, :] + i * m
        stacked_index[1, offset:offset + nnz] = sparse_tensor.indices()[1, :] + j * n
        
    return torch.sparse_coo_tensor(
        stacked_index,
        torch.tile(sparse_tensor.values(), (d * d,))
    ).coalesce()

這應該通過構造一個適當大小的空張量並填充索引，然后只是平鋪數據來解決二維張量的問題。