Dask：布尔索引原地修改引发异常行为问题排查

问题代码 #

import dask.array as da
i = da.full((10,10),fill_value=4)
c = da.ones(i.shape, dtype=bool)
c[:,0] = False # 增加多样性
i[c] |= da.full(i.shape, fill_value=2)
i[c]

预期输出 #

array([[6, 6, 4, 6, 6, 6, 6, 6, 6, 6],
       [6, 6, 4, 6, 6, 6, 6, 6, 6, 6],
       [6, 6, 4, 6, 6, 6, 6, 6, 6, 6],
       [6, 6, 4, 6, 6, 6, 6, 6, 6, 6],
       [6, 6, 4, 6, 6, 6, 6, 6, 6, 6],
       [6, 6, 4, 6, 6, 6, 6, 6, 6, 6],
       [6, 6, 4, 6, 6, 6, 6, 6, 6, 6],
       [6, 6, 4, 6, 6, 6, 6, 6, 6, 6],
       [6, 6, 4, 6, 6, 6, 6, 6, 6, 6],
       [6, 6, 4, 6, 6, 6, 6, 6, 6, 6]])

报错信息 #

ValueError: operands could not be broadcast together with shapes (nan,) (10, 10)

解决方法 #

Dask数组是不可变的，不支持NumPy式的掩码索引原地赋值操作，且i[c]返回的是形状延迟确定的数组，无法和原形状(10,10)的数组广播。正确的实现方式是用da.where函数进行条件替换，同时修正掩码数组的生成方式（Dask不允许直接修改数组切片）：

import dask.array as da

# 初始化原数组
i = da.full((10,10), fill_value=4)
# 生成掩码数组：第一列设为False，其余为True
c = da.where(da.arange(i.shape[1]) == 0, False, da.ones(i.shape, dtype=bool))
# 按掩码执行按位或操作：满足掩码的元素执行i | 2，否则保留原值
i = da.where(c, i | 2, i)

# 计算并输出结果
print(i.compute())

说明 #

1. 掩码数组生成：用da.where替代直接切片赋值，符合Dask不可变数组的特性。
2. 按位操作：通过da.where实现条件逻辑，对掩码为True的元素执行i | 2（4的二进制是100，2是010，按位或后为110即6），掩码为False的元素保持原数值4，最终得到预期结果。

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