WIP: For testing, optionally preserve 0-D arrays in operations #29067

seberg · 2025-05-27T10:13:02Z

This is a WIP, it works well enough for trying, but won't pass tests with environment variable set and has some smaller issues still.

We discussed with @mdhaber and others also to pursue removing the annoyance that 0-D arrays tend to be converted to scalars. This PR implements it, but hides it behind:

NUMPY_PRESERVE_0D_ARRAYS=1
np._get_preserve_0d_arrays()
np._set_preserve_0d_arrays()

All for testing and not thread/context-safe. We probably could do that but it isn't needed for testing. (I wouldn't want context use for indexing but it should be fine for anything else.)

This takes my "minimal change" approach, larger approaches were discussed, but the amount of test failures downstream didn't make it appealing on first sight. The changes are:

Any function/operator with 0-D arrays as input will return 0-D arrays. But functions that take only scalars will return scalar. (There are some subtlties, i.e. quantile(arr, q), only q dictates the result shape and thus only q is relevant.)
arr1d[0] is unchanged. arr[..., 0] or arr[0, ...] means that indexing already has a way to avoid getting scalars.
arr.sum(), i.e. axis=None is unchanged and returns scalars. But arr.sum(0) always returns arrays. The reason is that the first currently always returns a scalar, while the latter returns an (N-1) dimensional array or a scalar if 0-D.
- ufuncs actually default to axis=-1, I think, but typical reductions use axis=None.
- I think axis=None is special enough, although I admit it feels a bit different from the indexing ... use. (which may make sense as one lists all axis, while the other lists axes to be removed)

I would want to go with the above in a first merge, but we can discuss adding additional switches, e.g. for the behavior of reductions with axis=None. Indexing arr[()] could also be changed, but if we discuss this, we may want to start with adding something like arr.get_element(idx) first (because arr.item(), seems a bit awkward maybe).

Things to do here:

We need a CI run that passes with the environment variable
There is something wrong with masked arrays (have to investigate)
nanfuncs call np.asarray() early on and need to fixed if they are to return scalars for scalar inputs.
There are some other smaller bugs/test failures, at least in the env variable path right now.

This PR does "fix" that some functions return Python objects rather than NumPy scalars when inputs have dtype=object. I think this is OK, but it is a subtle change.

Signed-off-by: Sebastian Berg <sebastianb@nvidia.com>

Many remaining ones are around object arrays returning arguably better (less "NumPy") results, but probably not all :)

mhvk

@seberg - overall, this looks great, and I think it also helps to see why it is good not to be too strict about every result becoming arrays.

Also, it is nice to see the array converter that you implemented coming to the fore. That said, the old_scalar argument does not look nice - as asked inline, what exactly goes wrong without that?

mhvk · 2025-05-27T11:02:31Z

numpy/_core/src/multiarray/array_converter.c

@@ -404,7 +403,11 @@ static PyMethodDef array_converter_methods[] = {
        METH_FASTCALL | METH_KEYWORDS, NULL},
    {"wrap",
        (PyCFunction)array_converter_wrap,
-        METH_FASTCALL | METH_KEYWORDS, NULL},
+        METH_FASTCALL | METH_KEYWORDS,
+        "Apply array-wrap.  Supports `to_scalar=None/True/False` and "


I don't understand the docstring!

mhvk · 2025-05-27T11:11:16Z

numpy/_core/src/umath/ufunc_object.c

+
+    /*
+     * all_inputs_were_scalars is used to decide if 0-D results should be
+     * unpacked to a scalar.  But, `np.matmul(vector, vector)` should do this.


Also np.vecdot

mhvk · 2025-05-27T11:15:30Z

numpy/_core/src/umath/ufunc_object.c

+     * unpack.  (unless keepdims=True, since that behave like a normal ufunc)
+     * So we pretend inputs were scalars...
+     *
+     * TODO: We may need a way to customize that at some point.


For gufunc, whether an argument is scalar would seem to have the logical condition that there are no outer dimensions being iterated over. The effect would be the same as what you have here, but perhaps easier to describe?

Yeah, maybe that is a better way to describe "no outer dimensions in the result".

Unfortunately, that is also a nice way to explain why we don't like this behavior :). I.e. the whole point of this is that it means the number of outer dimensions (zero or more) has never an impact on the result type!
Unfortunately, unlike arr.sum(), vecdot(x, y) doesn't have a way to signal that we should return a scalar (by default).
Maybe out=... has to be the one way unfortunately... Or maybe vecdot should return arrays. Also, we could (and probably should!) make it so that if axes= is passed we never return a scalar (mirroring arr.sum(0)).

I see... I guess my initial sense would be to just always return arrays, at least for this trial... Do many tests break?

Note that axes=None (or axis=None) might still be a possibility - currently, that errors, so we can change it to mean "default axes, but return scalar if possible". Of course, that then prevents actually passing in meaningful axes and getting a scalar out. Also, it really is a hack.

Instead, maybe we should have an equivalent of out=... that makes explicit one does want a scalar if possible, say, out=() (or out=((), ()) for multiple outputs)?

Instead, maybe we should have an equivalent of out=... that makes explicit one does want a scalar if possible, say, out=() (or out=((), ()) for multiple outputs)

Maybe, I omitted multiple outputs for ... and I think it is better here as well. If we are going to go a step further here, I am wondering if even a return_scalar=True/False/None may not be better, though (and remove out=... again).

Maybe the main thing is if vec @ vec can change behavior. Because if we are willing to gamble on that, it may also be OK to just ask users to do np.matmul(vec, vec)[()] for starters. Or a more obvious arr.get_element()

EDIT: Part of me wonders if such a change might be OK to uncouple even. Also, we could allow gufuncs to opt out (always return arrays).

My sense would indeed be to uncouple, and just see if this causes big problems. I somewhat doubt it -- gufuncs are just not used as much and generally scalars are not really expected. (I definitely don't like the idea of a special keyword argument.)

mhvk · 2025-05-27T11:19:57Z

numpy/_core/tests/test_ufunc.py

+        assert type(np.sum(np.float32(2.5), axis=None)) is np.float32
+        assert type(np.max(np.float32(2.5), axis=None)) is np.float32
+        assert type(np.min(np.float32(2.5), axis=None)) is np.float32
+        # TODO: In a sense should return an array, but this is axis=0 is


Agree, arguably axis=0 on a 0-D array should just fail!

More relevantly here, are you sure you want to insist this returns a scalar? I guess just from being an operation on scalars only?

mhvk · 2025-05-27T11:24:51Z

numpy/lib/_arraysetops_impl.py


-    return _in1d(ar1, ar2, assume_unique, invert, kind=kind)
+    dt = conv.result_type()


dt does not seem to be used.

mhvk · 2025-05-27T11:34:25Z

numpy/lib/_nanfunctions_impl.py

-    a = np.asanyarray(a)
+    conv = _array_converter(a, q)
+    a, q_arr = conv.as_arrays(pyscalars="convert")
+
    if a.dtype.kind == "c":
        raise TypeError("a must be an array of real numbers")

    # Use dtype of array if possible (e.g., if q is a python int or float).
    if isinstance(q, (int, float)) and a.dtype.kind == "f":
        q = np.asanyarray(q, dtype=a.dtype)


Might as well change it to np.asarray here, like in the regular quantile.

numpy/polynomial/polyutils.py

mhvk · 2025-05-27T11:37:26Z

numpy/testing/_private/utils.py

@@ -1160,7 +1160,8 @@ def compare(x, y):
        if not issubdtype(z.dtype, number):
            z = z.astype(np.float64)  # handle object arrays

-        return z < 1.5 * 10.0**(-decimal)
+        # the float64 ensures at least double precision for the comparison.


Does this actually matter for this PR?

Yeah 3 linalg tests or so fail otherwise. But maybe should check whether the problem is really here or not.

mhvk · 2025-05-27T11:37:52Z

numpy/_core/multiarray.pyi

@@ -1289,3 +1289,6 @@ def nested_iters(
    casting: _CastingKind = ...,
    buffersize: SupportsIndex = ...,
 ) -> tuple[nditer, ...]: ...
+
+def _get_preserve_0d_arrays() -> bool: ...
+def _set_preserve_0d_arrays(state: bool, /) -> bool: ...


numpy/_core/src/multiarray/array_converter.c

mhvk · 2025-05-27T11:55:12Z

Mixing numpy scalars with array subclasses currently will lead to incorrect behavior (return arrays when it should scalars with the env variable).

Isn't subclasses always the right answer? Wouldn't __array_wrap__ have to deal with producing the scalar?

seberg · 2025-05-27T12:03:20Z

Isn't subclasses always the right answer? Wouldn't array_wrap have to deal with producing the scalar?

Kind of, although here we already ignored array-wrap, which should be OK. But basically, array-wrap would be passed return_scalar=True and that information is lost.
But, since __array_wrap__ is ignored, I think I can just add a if was_pyscalar: PyArray_Return() to these scalar functions. Before it was just unconditional, which was also wrong.

numpy/lib/_arraysetops_impl.py

mhvk · 2025-05-27T12:45:02Z

Mixing numpy scalars with array subclasses currently will lead to incorrect behavior (return arrays when it should scalars with the env variable).

Your earlier answer left me confused - when exactly does this happen? (only with preserve_0d_arrays, right?)

seberg · 2025-05-27T12:52:24Z

Mixing numpy scalars with array subclasses currently will lead to incorrect behavior (return arrays when it should scalars with the env variable).

Your earlier answer left me confused - when exactly does this happen? (only with preserve_0d_arrays, right?)

The initial comment was just wrong, sorry. We always potentially ignored __array_wrap__ behavior, but this had never anything to do with subclasses.
The original push didn't return the correct scalar for some numpy_scalar * object combinations in the new branch, I'll push a fix shortly.

Right now, more tests than I expected are failing with the env variable set, though.

numpy/__init__.pyi

numpy/_core/multiarray.pyi

numpy/_core/numeric.pyi

numpy/_core/src/multiarray/multiarraymodule.c

numpy/_core/src/npysort/x86-simd-sort

numpy/polynomial/polyutils.py

seberg · 2025-05-27T15:40:09Z

OK, things are a bit better now. Should pass without env var. The main problems currently:

nanfuncs use np.asarray(), they may all need the _array_converter treatment...
np.linalg test fails, because it returns array vs. scalar (which is also an array in the test) depending on ord= (I am not sure which version is correct, yet).
The masked scalar (and maybe other MaskedConstant) misbehave... It's an array and e.g. comparisons with it return masked arrays (and not the scalar).
And of course I didn't do anything to refine gufuncs yet.

numpy/__init__.pyi

numpy/_core/multiarray.pyi

numpy/_core/numeric.pyi

Co-authored-by: Joren Hammudoglu <jhammudoglu@gmail.com>

seberg added 9 commits May 27, 2025 10:14

WIP: Try to preserve scalars (squashed, as rebasing is too hard)

8b6a6a4

Try to use new logic for quantiles (not perfect yet, but better maybe?)

46fe997

Preserve scalar input in in1d

3d99b21

Use things to cast (move into helper could be nice here)

1ee236b

Signed-off-by: Sebastian Berg <sebastianb@nvidia.com>

A few more fixes after rebase

3db9414

silence some warnings

bea3c10

fixups to just reduce test failures a bit...

43501a7

Many remaining ones are around object arrays returning arguably better (less "NumPy") results, but probably not all :)

WIP: work towards env variable (python side needs looking into)

c3d46ff

undo scalar+subclass operator behavior for now

644a84a

github-actions bot added the 25 - WIP label May 27, 2025

seberg mentioned this pull request May 27, 2025

Is Array-In -> Scalar-Out OK? scientific-python/summit-2025#38

Open

mhvk reviewed May 27, 2025

View reviewed changes

seberg commented May 27, 2025

View reviewed changes

numpy/lib/_arraysetops_impl.py Outdated Show resolved Hide resolved

seberg changed the title ~~WIP: Preserve 0-D arrays in operations~~ WIP: For testing, optionally preserve 0-D arrays in operations May 27, 2025

seberg added 2 commits May 27, 2025 17:19

Fixups and cleanups

2b02180

undo style fix, because CI style fixer is wrong

362c9a3