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()
1条答案
按热度按时间tyky79it1#
我假设你试图平铺它,而不是扩展维度,虽然从问题中看有点不清楚,但通过操纵底层索引和数据Tensor就足够容易了。
对于2DTensor,这应该可以通过构造一个适当大小的空Tensor并填充索引,然后只是平铺数据来完成。