堆积稀疏矩阵

Question

我有以下功能的密集数组的现有代码，并希望支持稀疏矩阵以及

特别是，我有一个四维的tensor形状(k, k, n, n)，我把它部分地展平成matrix形状(k * n, k * n)，就像这样

# Roll the third dimension to the second position to get a tensor
# with shape (k, n, k, n)
rolled = np.rollaxis(tensor, 2, 1)
# Extract the dimensions
k, n = rolled.shape[:2]
# Reshape to partially flatten the array
matrix = rolled.reshape((k * n, k * n))

在稀疏设置中，我有一个kby k列表，其中每个元素是一个nby n稀疏矩阵。使用np.rollaxis不是一个选项，因为有些维度是稀疏的，有些维度是密集的。关于如何实现相同的行为，使我最终得到一个形状为(k * n, k * n)的稀疏矩阵，有什么想法吗

一种选择可能是使用coordinate format并分别设置matrix的每个元素，但这看起来相当麻烦且容易出错