Numba @jit(nopython=True)下掩码除法求均值报错解决方法

Numba nopython=True 模式对np.divide的where参数兼容性存在缺陷：

• 低版本Numba会直接抛出编译错误，提示不支持该参数传入
• 高版本即使可编译通过，mask未覆盖的数组位置不会被自动初始化，存在脏数据，后续计算均值会得到错误结果

无需手写任何循环，直接使用Numba nopython模式原生支持的NumPy接口组合即可实现等价逻辑，代码如下：

import numpy as np
from numba import jit

A = np.array([[2,2,2],[1,0,0],[1,2,1]], dtype=np.float32)
B = np.array([[2,0,2],[0,1,0],[1,2,1]],dtype=np.float32)
C = np.array([[2,0,1],[0,1,0],[1,1,2]],dtype=np.float32)

@jit(nopython=True)
def test(a,b,c):
    denom = a + b
    # 替换零分母为1，避免除零生成inf或触发警告
    safe_denom = np.where(denom > 0, denom, 1.0)
    div = c / safe_denom
    # 原分母为0的无效位置标记为nan
    div[denom <= 0] = np.nan
    # 按行计算均值，自动跳过nan位置
    result = np.nanmean(div, axis=1)
    return result

test_res = test(A,B,C)

逻辑说明 #

• 用到的np.where、数组索引赋值、np.nanmean均属于Numba nopython模式原生支持的接口，可正常编译运行
• 实现逻辑和原代码完全等价：分母大于0的位置正常做除法，分母为0的位置不参与均值计算
• 运行结果和原生NumPy执行原代码的输出完全一致

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