Check if matrix contains only specific values in numpy

Minimal fix

There are two reasons why the code failed:

1. You cannot compare to NaN with m0 == np.nan because floating point rules (IEEE-754) dictate that comparisons with NaN are always false. You have to use np.isnan instead
2. Your binary logic is false and should use | instead of &

The minimal fix is

print(np.all((m0 == 0) | np.isnan(m0)))

Performance

However, there is a performance issue: np.unique and np.uniqueall create sorted output. It may optimize by using a hash table (taking O(N) time and O(M) space to scan N entries with M unique values) and only sort the resulting unique values in O(M log M) time but even that gives you a worst-case performance of O(N log N) if all values are unique.

# best case input because only one unique value
and np.all can stop after the first value if it
uses short-circuit logic
m = np.ones((1000, 1000))
Version 1, original code
np.all(np.uniqueall(m).values == 0)
Version 2, don't compute unused outputs in uniqueall

np.all(np.uniquevalues(m) == 0)
Version 3, compatible with Numpy 1.x
np.all(np.unique(m) == 0)
Version 4, optimal, simple
np.all(m == 0)

In Numpy 1.26, only version 3 and 4 work. Among those, version 3 is 25 times slower than version 4. In Numpy 2.4.6, version 1 is 44 times slower while version 2 is 7 times slower, same as version 3.

Checking many values

If you want to check multiple values, numpy comes with a suitable function: np.isin. So if you want to check, let's say 0, 2, 4, 6, you can write

(np.isin(m, [0, 2, 4, 6]) | np.isnan(m)).all()