使用map_blocks和reduce运行compute时触发TypeError的排查求助

问题描述 #

代码逻辑为：对数组按变量组处理，同一组内通过map_blocks执行卷积运算，单元素组直接保留二维数组；之后执行da.argmax并索引vals数组，最终拼接结果。执行loc.compute()时触发如下错误：

File c:\Users\x\lib\site-packages\dask\array\slicing.py:990, in check_index(axis, ind, dimension)
    987 elif ind is None:
    988     return
--> 990 elif ind >= dimension or ind < -dimension:
    991     raise IndexError(
    992         f"Index {ind} is out of bounds for axis {axis} with size {dimension}"
    993     )

TypeError: '>=' not supported between instances of 'str' and 'int'

已尝试scheduler='single-threaded'调试，简化代码后定位错误源自da.argmax(res, axis=1)环节，且确认数组无空值/无穷值。

错误原因分析 #

核心问题是调用da.map_blocks时未手动指定输出的chunks参数：

• 卷积函数中，da.fft.irfft(a * b, n = size)强制指定了输出长度为size，导致处理后每个chunk的数组形状与输入chunk的形状不一致（输入axis=1长度为1000，输出变为10000）。
• Dask默认会根据输入数组的形状推断输出形状，当实际输出形状与推断结果不符时，会导致元数据（如维度大小、块结构）错误，后续da.argmax和索引操作时，内部检查环节出现类型不匹配（将错误的元数据类型当成字符串与整数比较）。

解决方法 #

1. 为map_blocks指定输出chunks参数 #

在调用da.map_blocks时，明确指定输出数组的块结构，确保Dask能正确识别输出形状：

以简化代码为例，修改res = da.map_blocks(...)为：

# 输出形状为 (size, 10000)，块结构对应输入的第一个维度块大小，第二个维度为10000（单块）
output_chunks = (test[:,0].chunks[0], (size,))
res = da.map_blocks(simple_convolve, test[:,0], test[:,1], dtype='float32', chunks=output_chunks)

在原始代码中，修改reduce内的map_blocks调用：

# 确定输出的块结构：第一个维度与输入的chunk一致，第二个维度为size
output_chunks = (test_list[0].chunks[0], (size,))
temp = reduce(lambda x, y: da.map_blocks(custom_convolve, x, y, dtype='float32', chunks=output_chunks), test_list)

2. 优化卷积函数（可选） #

原函数的循环处理可改用向量化操作替代，简化代码并提升效率：

def custom_convolve(x,y):
    # 直接对整个数组做FFT，无需循环
    a = da.fft.rfft(x, axis=1)
    b = da.fft.rfft(y, axis=1)
    conv_res = da.fft.irfft(a * b, n=size, axis=1)
    return conv_res

优化后仍需手动指定chunks参数，确保形状推断正确。

3. 验证索引合法性 #

确认da.argmax(temp, axis=1)返回的索引值范围在vals的长度内：

• 如果temp的axis=1长度是size，且vals的长度也为size，则索引合法；若两者长度不一致，需调整irfft的n参数或vals的长度。

额外说明 #

关于da.isnan(res).sum().compute().compute()：
第一次compute()已将Dask数组转换为numpy数组，第二次compute()完全多余，正确写法为da.isnan(res).sum().compute()，直接得到numpy标量结果。

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