Creating block matrix from coo_matrix

Question

I have a torch.sparse_coo of size nxn.

How can I expand it into a ndxnd torch.sparse_coo matrix, where each entry is replaced by a dxd matrix of the same value?

I would like to do this without ever converting into dense.

Thank you!

Answer 1

I assume you're trying to tile it and not expand the dimensionality, although it's a bit unclear from the question. It'll be easy enough to do by manipulating the underlying indices and data tensors.

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()

This should do the trick for a 2D tensor by constructing an empty tensor of appropriate size and filling out the indices, then just tiling the data.