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

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:

Checks for equality to zero
If all zeros, it skips things
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:

Create an empty array
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.

@hmaarrfk

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.

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").

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(..) to all_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 (think add_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.

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....

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.

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?

Specifically, my understanding is that

np.all(arr == 7)

would decompose into:

A memory allocation of O(arr.size) for the result of arr == 7
Memory read from RAM (all cache misses for "large" arrays) for the contents of arr
Memory writes, to RAM (all causing cache evictions) for the contents of arr == 7
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.

is there a path for numpy to optimize operations like

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

Thanks or the discussion:

For me the discussion has led to two options:

Option 1: No new API: Numpy must figure out how to optimize array_equal. JITs must "inspect" things like strides to optimize things.

a = np.zeros((100, 1024, 1024), dtype='float64')
def array_equal_to_scalar(array, scalar):
    s = np.broadcast_to(np.array(scalar, dtype=array.dtype), array.shape)
    return np.array_equal(array, s)
array_equal_to_scalar(a, 7)

Option 2: Create add to the API to help guide people to this optimization

a = np.zeros((100, 1024, 1024), dtype='float64')
np.array_equal_to_scalar(array, 7)

I personally think that the whole broadcast_to trick is so niche it often looks like bad code.

Just reading it, I feel compelled to, in 2 years to write code like:

# This code will help trigger a fastpath in numpy that
# detects that the strides of the "seven" array array
# are zero, and just "compare to a scalar"
seven = np.broadcast_to(np.array(7, dtype=array.dtype), array.shape)

IMO: These 3 lines of comments, just feel "wrong" to me.

Navigation Menu

Repository navigation

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

Metadata

Assignees

Labels

Type

Projects

Milestone

Relationships

Development

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

Description

Activity

rgommers commented on Dec 22, 2025

hmaarrfk commented on Dec 22, 2025

rgommers commented on Dec 23, 2025

hmaarrfk commented on Dec 23, 2025

rgommers commented on Dec 23, 2025

hmaarrfk commented on Dec 23, 2025

hmaarrfk commented on Dec 23, 2025

rgommers commented on Dec 24, 2025

hmaarrfk commented on Dec 26, 2025

Metadata

Metadata

Assignees

Labels

Type

Projects

Milestone

Relationships

Development

Issue actions