Dask数组与scipy dia_matrix运算报错：'dia_matrix'不可下标访问

问题重现 #

import dask.array as da
from scipy.sparse import identity

I = identity(35344)
A = da.random.randint(0, high=5, size=(35344,35344))
obj_value = da.linalg.norm(A + I, ord="fro")
obj_value = float(obj_value)

报错信息 #

TypeError                                 Traceback (most recent call last)
c:\Users\samke\OneDrive - University of Kansas\Documents\KU\Research\Ashourvan Lab\Shared Papers\code\unfinished code\structure_function\test.ipynb Cell 4 in ()
      6 A = da.random.randint(0, high=5, size=(35344,35344))
      8 obj_value = da.linalg.norm(A + I, ord="fro")
----> 9 print(float(obj_value))

...（中间栈信息省略）

File c:\Users\samke\AppData\Local\Programs\Python\Python310\lib\site-packages\dask\array\core.py:118, in getter(a, b, asarray, lock)
    116     lock.acquire()
    117 try:
---> 118     c = a[b]
    119     # Below we special-case `np.matrix` to force a conversion to
    120     # `np.ndarray` and preserve original Dask behavior for `getter`,
    121     # as for all purposes `np.matrix` is array-like and thus
    122     # `is_arraylike` evaluates to `True` in that case.
    123     if asarray and (not is_arraylike(c) or isinstance(c, np.matrix)):

TypeError: 'dia_matrix' object is not subscriptable

报错原因 #

Dask数组与scipy原生稀疏矩阵（dia_matrix）混合运算时，Dask内部任务会尝试对dia_matrix执行下标访问（a[b]），但scipy的dia_matrix不支持这种操作。numpy处理小矩阵时正常，是因为numpy会自动将稀疏矩阵转为稠密数组，不会触发该问题。

解决方法 #

方法1：使用Dask稀疏数组（推荐，节省内存） #

先安装dask-sparse库，将scipy稀疏矩阵转为Dask稀疏数组后再运算：

import dask.array as da
from scipy.sparse import identity
import dask_sparse as ds

I = identity(35344)
# 转换为Dask稀疏数组
I_ds = ds.from_scipy(I)
A = da.random.randint(0, high=5, size=(35344,35344))
obj_value = da.linalg.norm(A + I_ds, ord="fro")
obj_value = float(obj_value)

方法2：转为稠密Dask数组（内存足够时使用） #

将scipy稀疏矩阵转为稠密numpy数组，再分块转为Dask数组：

import dask.array as da
from scipy.sparse import identity

I = identity(35344).toarray()
# 分块转为Dask数组，可根据内存调整chunks大小
I_da = da.from_array(I, chunks=(1000, 1000))
A = da.random.randint(0, high=5, size=(35344,35344))
obj_value = da.linalg.norm(A + I_da, ord="fro")
obj_value = float(obj_value)

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