使用Numba实现numpy.nanmean时遭遇TypingError问题求助

用户尝试用Numba编写简化版的numpy.nanmean函数，代码如下：

from numba import jit, prange
import numpy as np

@jit(nopython=True)
def nanmeanMY(a, axis=None):
    if a.ndim>1:
        ncols = a.shape[1]
        nrows = a.shape[0]
        a = a.T.flatten()
        res = np.zeros(ncols)
        for i in prange(ncols):
            col_no_nan = a[i*nrows:(i+1)*nrows]
            res[i] = np.mean(col_no_nan[~np.isnan(col_no_nan)])
        return res
    else:
        return np.mean(a[~np.isnan(a)])

预期功能：判断输入是向量还是矩阵，矩阵时返回各列的均值。使用测试矩阵：

X = np.array([[1,2], [3,4]])
nanmeanMY(X)

运行后抛出如下错误：

Traceback (most recent call last):

  Cell In[157], line 1
    nanmeanMY(a)

  File ~\anaconda3\Lib\site-packages\numba\core\dispatcher.py:468 in _compile_for_args
    error_rewrite(e, 'typing')

  File ~\anaconda3\Lib\site-packages\numba\core\dispatcher.py:409 in error_rewrite
    raise e.with_traceback(None)

TypingError: No implementation of function Function(<built-in function getitem>) found for signature:
 
getitem(array(int32, 2d, C), array(bool, 2d, C))
 
There are 22 candidate implementations:
      - Of which 20 did not match due to:
      Overload of function 'getitem': File: <numerous>: Line N/A.
        With argument(s): '(array(int32, 2d, C), array(bool, 2d, C))':
       No match.
      - Of which 2 did not match due to:
      Overload in function 'GetItemBuffer.generic': File: numba\core\typing\arraydecl.py: Line 209.
        With argument(s): '(array(int32, 2d, C), array(bool, 2d, C))':
       Rejected as the implementation raised a specific error:
         NumbaTypeError: Multi-dimensional indices are not supported.
  raised from C:\Users\****\anaconda3\Lib\site-packages\numba\core\typing\arraydecl.py:89

During: typing of intrinsic-call at C:\Users\****\AppData\Local\Temp\ipykernel_10432\1652358289.py (22)

• Numba的nopython模式不支持使用多维布尔数组作为索引来切片多维数组，错误提示里的NumbaTypeError: Multi-dimensional indices are not supported已经明确指出这一点。
• 虽然运行时输入二维数组会进入if a.ndim>1分支，但Numba在编译阶段会对所有分支的代码做类型检查。当输入为二维数组时，else分支里的a[~np.isnan(a)]会被解析为用二维布尔数组索引二维数组，直接触发类型推断错误。
• 额外注意：即使解决了这个问题，代码还存在潜在风险——当某列全为NaN时，col_no_nan[~np.isnan(col_no_nan)]会变成空数组，Numba对np.mean处理空数组的逻辑和NumPy可能存在差异，需要额外处理。

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