Python 3.x 使用嵌套的np.less对np.all进行短路，以便在numpy中进行大数组比较

在我当前的代码（参见MWE）中，我有一个瓶颈，在这里我对大型2D数组执行

np.all

和嵌套

np.less

。我知道，如果

np.less

中有一个

false

值，我们可以停止检查，因为索引中的其余值代码将计算为

false

（因为我是并且将给定维度的单个索引中的所有值合并在一起）

使用numba或numpy有没有办法利用这种“早期退出/短路”条件在计算中产生有意义的加速

MWE中倒数第二行是我要加速的。请注意，

N

和

M

可能非常大，但实际上只有很少的比较结果是

true

import numpy as np

N = 10000
M = 10 # Reduced to small value to show that sometimes the comparisons evaluate to 'True'

array = np.random.uniform(low=0.0, high=10.0, size=(N, M))
comparison_array = np.random.uniform(low=0.0, high=10.0, size=(M))

# Can we apply an early exit condition on this?
mask = np.all(np.less(array, comparison_array), axis=-1)

print(f"Number of 'True' comparisons: {np.sum(mask)}")

---

以下是

numba

版本，开发到足以工作，但不一定经过优化：

@numba.njit
def foo(arr, carr):
    N, M = arr.shape
    mask = np.ones(N, dtype=np.bool_)
    for i in range(N):
        for j in range(M):
            if arr[i,j]>=carr[j]:
                mask[i]=False
                break
    return mask

测试：

In [178]: np.sum(foo(array, comparison_array))
Out[178]: 2
In [179]: np.sum(np.all(np.less(array, comparison_array), axis=1))
Out[179]: 2

时间：

In [180]: timeit np.sum(foo(array, comparison_array))
155 µs ± 6.36 µs per loop (mean ± std. dev. of 7 runs, 10000 loops each)
In [181]: timeit np.sum(np.all(np.less(array, comparison_array), axis=1))
451 µs ± 5.19 µs per loop (mean ± std. dev. of 7 runs, 1000 loops each)

---

这是一个不错的进步。

哇，这是一个很大的进步！我应该关注哪些事情来构建您的解决方案以进一步优化它？