如何从torch.sparse_coo矩阵构建分块矩阵（无需转稠密）

要实现将n×n的torch.sparse_coo_tensor扩展为nd×nd的稀疏矩阵，且每个原元素替换为同值的d×d块，核心是扩展稀疏坐标并重复对应值，全程无需转换为稠密矩阵。以下是具体实现步骤：

核心思路 #

原稀疏矩阵的每个非零元素位于坐标(i,j)，对应扩展后矩阵的(i*d + k, j*d + l)位置（其中k,l ∈ [0, d-1]），这些位置的值均与原元素相同。我们只需生成所有扩展后的坐标，并将原值重复d²次即可。

代码实现与解释 #

import torch

# 1. 创建示例n×n稀疏COO矩阵
n = 2
d = 2
orig_indices = torch.tensor([[0, 1, 1], [1, 0, 1]])  # 非零元素坐标
orig_values = torch.tensor([1.0, 2.0, 3.0])          # 对应值
orig_sparse = torch.sparse_coo_tensor(orig_indices, orig_values, size=(n, n))

# 2. 提取原矩阵的坐标和值
orig_i = orig_indices[0]  # 行坐标，shape (N,)，N为非零元素数量
orig_j = orig_indices[1]  # 列坐标，shape (N,)
orig_vals = orig_sparse.values()

# 3. 生成d×d块的偏移量
k = torch.arange(d)  # 行偏移：0,1,...,d-1
l = torch.arange(d)  # 列偏移：0,1,...,d-1

# 4. 计算扩展后的坐标
# 每个原行坐标i对应d个行位置：i*d + 0, i*d +1, ..., i*d +d-1，每个位置重复d次（对应所有列偏移）
new_i = (orig_i.unsqueeze(1) * d + k).repeat_interleave(d, dim=1).flatten()
# 每个原列坐标j对应d个列位置：j*d +0, j*d +1, ..., j*d +d-1，每个位置被每个行偏移对应一次
new_j = (orig_j.unsqueeze(1) * d + l).repeat(1, d).flatten()
new_indices = torch.stack([new_i, new_j])

# 5. 生成扩展后的值：每个原值重复d²次
new_vals = orig_vals.repeat_interleave(d*d)

# 6. 创建扩展后的nd×nd稀疏COO矩阵
new_sparse = torch.sparse_coo_tensor(new_indices, new_vals, size=(n*d, n*d))

# 验证（可选，仅用于确认结果）
print("原矩阵稠密形式：")
print(orig_sparse.to_dense())
print("\n扩展后矩阵稠密形式：")
print(new_sparse.to_dense())

关键细节说明 #

• 坐标扩展：通过广播和重复操作，避免了循环遍历非零元素，保证了效率；
• 值的重复：使用repeat_interleave将每个原值重复d²次，与扩展后的坐标一一对应；
• 全程稀疏操作：所有步骤均基于原稀疏矩阵的indices和values张量操作，未转换为稠密矩阵，适合大尺寸矩阵场景。

内容的提问来源于stack exchange，提问作者Neurobro