开启parallel=True时Numba编译报‘key already in dictionary’错误求助

使用Anaconda提供的Python 3.11版本编写了如下函数，当为@njit装饰器设置parallel=True选项时，Numba无法完成编译：

@njit(nopython=True, nogil=True, cache=True, parallel=True)
def corrected_jacobian_numba_icecore(age_top_sigma,
                             accu, tau, lid, dens, dens_firn,
                             depth, depth_inter, depth_mid,
                             age, airage, age_model, airage_model,
                             ice_equiv_depth,
                             agedens, icelayerthick,
                             corr_a, corr_tau, corr_lid,
                             chol_a, chol_tau, chol_lid,
                             sigmap_corr_a, sigmap_corr_tau, sigmap_corr_lid,
                             corr_a_age, corr_tau_depth, corr_lid_age):

    age_jac = np.zeros((1+len(corr_a)+len(corr_tau)+len(corr_lid), len(age)))
    age_jac[0, :] = age_top_sigma * np.ones(len(age))       
    airage_jac = np.zeros((1+len(corr_a)+len(corr_tau)+len(corr_lid), len(airage)))
    airage_jac[0, :] = age_top_sigma * np.ones(len(airage))       
    delta_depth_jac = np.zeros((1+len(corr_a)+len(corr_tau)+len(corr_lid), len(depth)))
#    delta_depth[0, :] = age_top_sigma * np.zeros(len(depth))  # useless

    for i in prange(len(corr_a)):

        corr_a_vec = np.zeros(len(corr_a))
        corr_a_vec[i] = 1.
    #Accu
        corr_vec = np.dot(chol_a, corr_a_vec)*sigmap_corr_a
        toto = np.interp((age_model[:-1]+age_model[1:])/2,
                                  corr_a_age, corr_vec)
        agedens_vec = - toto * agedens

    #Ice age
        age_vec = np.cumsum(np.concatenate((np.array([0]), depth_inter*agedens_vec)))
        age_jac[1+i,:] = age_vec

    #Air age
        airage_vec = np.interp(ice_equiv_depth, depth, age_vec)
        airage_jac[1+i, :] = airage_vec
        # delta_depth_vec = np.zeros_like(depth)
        # delta_depth_jac[1+i, :] = delta_depth_vec
            

        
    for i in prange(len(corr_tau)):
                        
        corr_tau_vec = np.zeros(len(corr_tau))
        corr_tau_vec[i] = 1.
        corr_vec = np.dot(chol_tau, corr_tau_vec)*sigmap_corr_tau
        tata = np.interp(depth_mid, corr_tau_depth, corr_vec)
        agedens_vec = -tata * agedens                
        age_vec = np.cumsum(np.concatenate((np.array([0]), depth_inter*agedens_vec)))
        age_jac[1+len(corr_a)+i, :] = age_vec
        
        thin_vec = -tata * dens/tau
        udepth_vec = np.cumsum(np.concatenate((np.array([0]), depth_inter*thin_vec)))
        delta_depth_vec = - np.interp(ice_equiv_depth, depth_mid,
                                      tau/dens) * (udepth_vec - \
                                    np.interp(ice_equiv_depth, depth, udepth_vec))
        airage_vec = np.interp(ice_equiv_depth, depth, age_vec) \
                        - np.interp(ice_equiv_depth, depth_mid, agedens) * \
                        delta_depth_vec
        airage_jac[1+len(corr_a)+i, :] = airage_vec
        delta_depth_jac[1+len(corr_a)+i, :] = delta_depth_vec

        #To be continued...                

    for i in prange(len(corr_lid)):

        age_vec = np.zeros_like(depth)
        age_jac[1+len(corr_a)+len(corr_tau)+i, :] = age_vec

        corr_lid_vec = np.zeros(len(corr_lid))
        corr_lid_vec[i] = 1.
        corr_vec = np.dot(chol_lid, corr_lid_vec)*sigmap_corr_lid
        lid_vec = np.interp(airage_model, corr_lid_age, corr_vec) * lid
        delta_depth_vec = dens_firn * lid_vec * \
                            np.interp(ice_equiv_depth, depth_mid, 
                                      tau/dens)
        airage_vec = - np.interp(ice_equiv_depth, depth_mid, 
                                  agedens) * delta_depth_vec
        airage_jac[1+len(corr_a)+len(corr_tau)+i, :] = airage_vec
        delta_depth_jac[1+len(corr_a)+len(corr_tau)+i, :] = delta_depth_vec

    return airage_jac, delta_depth_jac, age_jac

编译时出现的错误信息：

AssertionError: Failed in nopython mode pipeline (step: nopython rewrites)
key already in dictionary: '$np_typ_var.1'

解决方法 #

1. 升级Numba版本 #

这个错误是Numba的已知类型推断bug，在Numba 0.57及以上的稳定版本中已修复。可以通过以下命令升级：

conda update numba

或使用pip：

pip install --upgrade numba

2. 优化循环内临时数组的创建 #

在prange循环内重复创建结构相同的数组（如corr_a_vec、corr_tau_vec）会干扰Numba的并行类型推断。可将数组预分配到循环外部，循环内仅重置值：

示例修改第一个循环：

# 预分配数组到循环外
corr_a_vec = np.zeros(len(corr_a))
for i in prange(len(corr_a)):
    corr_a_vec[:] = 0.  # 重置为0
    corr_a_vec[i] = 1.
    # 后续逻辑保持不变
    corr_vec = np.dot(chol_a, corr_a_vec)*sigmap_corr_a
    toto = np.interp((age_model[:-1]+age_model[1:])/2, corr_a_age, corr_vec)
    agedens_vec = - toto * agedens
    # ...

同理修改另外两个prange循环的临时数组创建逻辑。

3. 暂时禁用cache选项 #

cache=True与parallel=True同时启用时，可能触发编译缓存冲突。可先去掉cache=True编译函数，确认通过后再重新开启缓存。

4. 拆分复杂的interp调用 #

循环内多次调用np.interp可能导致类型推断冲突，可将重复的interp逻辑提取为单独的@njit函数：

@njit
def interp_helper(x, xp, fp):
    return np.interp(x, xp, fp)

然后在主函数中替换所有np.interp为interp_helper调用。

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