运行包含map_blocks和reduce的计算时出现类型错误

Question

我很难诊断错误的原因。我的代码涉及对某些数组运行卷积（使用

map_blocks

）（如果它们属于同一组变量），否则只记录 2 维数组。然后，我执行

argmax

操作并将结果添加到列表中，然后将其连接起来。

我尝试使用

scheduler='single-threaded'

参数运行计算来帮助调试，但我仍然无法看到错误的原因。

import dask.array as da
from functools import reduce
import numpy as np

size = 100000
vals = da.linspace(0, 1, size)
nvars = 12
test = da.random.uniform(low=0, high=1, size=(100000, nvars, size), chunks=(100, nvars, size))

# number of total unique items corresponds to nvars
var_lookup = {
        'a': [0],
        'b':
        [0, 1, 2],
        'c': [0, 1],
        'd': [0, 1, 2],
        'e': [0],
        'f': [0],
        'g': [0],
    }

# Iterates over all 0 dimension coordinates
# and convolves relevant values from x and y
def custom_convolve(x,y):
    temp_lst = []
    for i in range(x.shape[0]):
        a = da.fft.rfft(x[i])
        b = da.fft.rfft(y[i])
        conv_res = da.fft.irfft(a * b, n = size)
        temp_lst.append(conv_res)
    res = da.stack(temp_lst, axis=0)
    return res

n_groups = len(var_lookup.keys())

counter = 0
group_cols = []
for i in var_lookup.keys():
    grp = var_lookup[i]
    # if group consists of 1 value, then just record that 2-dim array
    if len(grp)==1:
        temp =  test[:,counter,:]
        counter += 1
    else:
        test_list = []
        for _ in var_lookup[i]:
            test_list.append(test[:, counter, :])
            counter += 1
        temp = reduce(lambda x, y: da.map_blocks(custom_convolve, x, y, dtype='float32'), test_list)

    res = vals[da.argmax(temp, axis=1)]

    group_cols.append(res)

loc = da.stack(group_cols, axis=1)

运行计算时出错：

res = 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'

也许

reduce

函数与

map_blocks

相结合导致了问题？