[RFC]: add APIs for array equality to a scalar

[hmaarrfk]

Today, there are a few cases where people might want to check equality to a scalar.

import numpy as np

is_all_zero = np.all(arr == 0)

However, this code is terribly inefficient for out of memory array.

It can force them to go through the entire memory.

I think that we can do much better to define an operation that will check for

• All equal
• All not equal

and those can be implemented in more "streamed" fashion, that allow the underlying implementation to not create a full boolean array for arr == 0.

Second, I would like to propose an API for "likely not equal", where the strick inequality is not guaranteed.

For data compression, we might just be intrested in learning if the dataset is worth compressing or not:
Zarr for example does this:
zarr-developers/zarr-python#3627

For example, consider the task to compress a 800MB array.
Zarr today:

1. Checks for equality to zero
2. If all zeros, it skips things
3. if not zeros, it compresses things

But if you have an out of memory data array, it may be hard to guarantee that things are not all zero, but the check likely doesn't matter, since with modern compression algorithm all zeros can be efficiently compressed.

>>> import numpy as np
>>> import numcodecs
>>> a = np.zeros((100, 1024, 1024), dtype='float64')
>>> len(numcodecs.blosc.compress(a, b'zstd', 7))
67216

On my computer, the equality check is roughly the same order of magnitude as the compression itself:

In [14]: %time numcodecs.blosc.compress(a, b'zstd', 7);
CPU times: user 2.04 s, sys: 903 μs, total: 2.04 s
Wall time: 266 ms

In [15]: %time np.all(a==0 );
CPU times: user 21.4 ms, sys: 109 ms, total: 130 ms
Wall time: 140 ms

so if zarr misses one of my images, because I think it has non-zero element, its likely not the end of the world in terms of system performance, but if I have to "guarantee that the images are non zero" that can be a costly operation that can't be done through my implementation easily "without checking every single element".

For not, without these APIs, today, zarr does something like:

1. Create an empty array
2. Check for equality with that empty array

---

Alternative:

I could likely check numpy metadata like strides and create my own fastpath for this:

>>> zero = np.zeros(1)
>>> a, zero = np.broadcast_arrays(a, zero)
>>> zero.strides
(0, 0, 0)

but this seems more like a hack and would be implicitely redefining "equality operation" to "likely equal" which isn't correct.

[rgommers]

Thanks for the ideas @hmaarrfk

However, this code is terribly inefficient for out of memory array.

It can force them to go through the entire memory.

count_nonzero(arr) and any(arr) (and not any(arr)) can both short-circuit, so should be more efficient than np.all(arr == 0).

Second, I would like to propose an API for "likely not equal", where the strick inequality is not guaranteed.

That doesn't seem feasible to standardize, because (a) it doesn't exist in current array libraries, and (b) it probably doesn't make sense for most libraries (e.g., why would numpy want to add this?).

I could likely check numpy metadata like strides and create my own fastpath for this:

That doesn't seem worth doing. Serializing a just-allocated array with all zeros is a very special case, and doesn't really occur in well-written code.

[hmaarrfk]

count_nonzero(arr) and any(arr) (and not any(arr)) can both short-circuit, so should be more efficient than np.all(arr == 0).

These are actually pretty good tips... I should have thought of these, or tried to copy this message to AI to suggest them to the zarr project.

In trying to implement it, I realize that zarr doesn't only only check for "equality to zero" but rather equality to the fill value, which can be non-zero.

To do so, they somewhat defined their own function all_equal.

I'm happy to discuss the inclusion of an API where we check equality to a scalar, or an "other", and not consider the other "proposals" i made (including the "likely not").

[rgommers]

Note that that's pretty unlikely to fly. Functions like that may in some cases make sense for eager libraries, but we already have all and equal in the standard, and JIT compilers can transparently optimize all(equal(..)toall_equal(...)so there's no benefit to a separate function. And there are an endless amount of composed operations that are faster in eager mode (thinkadd_mul`, etc.) - NumPy et al don't have those because the tradeoffs of increased API surface vs. performance gain on the operation in quesiton aren't good.

[hmaarrfk]

My understanding was that JITs were still complicated to use in Windows.

Personally, i tend to avoid them due to the "warmup costs".

I haven't used them in a while, maybe numba got alot of better, and the "caching" of the compiled results is better today, but I think a real concern is startup time for short lived scripts.

AOT compilation is still valuable in 2026 IMO....

[rgommers]

You're all correct there, there's an overhead in using JITs. But we still have to deal with design that work for JIT/AOT/eager, so we're very reluctant to add eager-only APIs. There has to be a really compelling case, and performance of a composed operation isn't that.

[hmaarrfk]

Perhaps, i'm missing something, is there a path for numpy to optimize operations like

np.all(arr == 0)

without allocating an entire other array of size arr.size bytes?

[hmaarrfk]

Specifically, my understanding is that

np.all(arr == 7)

would decompose into:

1. A memory allocation of O(arr.size) for the result of arr == 7
2. Memory read from RAM (all cache misses for "large" arrays) for the contents of arr
3. Memory writes, to RAM (all causing cache evictions) for the contents of arr == 7
4. A memory read of the order of arr.size to establish that all the elements are in fact 7.

In the above decomposition, 1, 3, and 4 are all unnecessary for this kind of "is my array really a constant" kind of question.

[rgommers]

is there a path for numpy to optimize operations like

There isn't for numpy, your analysis looks correct.