Repository: numpy/numpy
Pull request: BUG: Remove unnecessary copying and casting from out array in choose
Resolves #22237.

This PR does not copy the existing `out` array in `np.choose`, instead creating the new output array from scratch and copying it back manually.

Thank you for reviewing!


Problem statement:
BUG: `choose` does an unnecessary cast of the initial `out` values
### Describe the issue:

Choose (I think) always overwrites all values in the output.  It is thus unnecessary to cast a given `out` to the resulting dtype (if a cast is necesssary).

This can lead to spurious warnings as given in the below example

### Reproduce the code example:

```python
import numpy as np

a = np.arange(3)
b = np.ones((3, 3), dtype=np.int64)

out = np.array([np.nan, np.nan, np.nan])

np.choose(a, b, out=out)
```


### Error message:

```shell
/home/sebastian/forks/numpy/build/testenv/lib/python3.9/site-packages/numpy/core/fromnumeric.py:57: RuntimeWarning: invalid value encountered in cast
  return bound(*args, **kwds)
```


### NumPy/Python version information:

All versions up to 1.24-dev will do the unnecessary cast, but only the current dev gives the spurious warning.
(We have a typing test that could run into this randomly, although I am changing it in a different PR)

### Context for the issue:

_No response_

Hints from issue discussion:
Hmm, I looked at this one for a bit, it definitely wasn't as simple as naively preserving the `dtype` of `out` in `obj` that gets assigned to. Such a strategy did allow me to pass the example above without a warning, but some other tests were really sensitive to a variety of issues related to memory overlap and some nasty tests where `byteorder` was different between some of the arrays involved. I tried detecting that and using `PyArray_Byteswap` but it didn't seem to help much.

I started messing around with `PyArray_Descr` tailored to specific purposes and it went about as well as you might expect. Maybe I'll circle back to this in the sprint tomorrow, but I definitely ended up in a situation where patching the immediate issue made other parts of the testsuite pretty unhappy.
I think the problem here isn't the new array/dtype.  We need the new array with writebackifcopy for casting.  But `PyArray_FromArray` copies the old array when necessary, while what we actually somewhat need is a new empty one (which we copy when we are done).

Not sure whats the best way to achieve this, maybe it needs a manual check or a new helper, unless some `out` values need to be preserved, that seems always better in principle.  If they do need to be preserved, things are a bit funny anyway, because allowing casts that are not safe is awkward.

Commit message: BUG: Add defensive comparison instead of equals for failure check

Patch to review:
diff --git a/numpy/_core/src/multiarray/item_selection.c b/numpy/_core/src/multiarray/item_selection.c
index a6c92dc00edc..607045444cbb 100644
--- a/numpy/_core/src/multiarray/item_selection.c
+++ b/numpy/_core/src/multiarray/item_selection.c
@@ -1019,23 +1019,8 @@ PyArray_Choose(PyArrayObject *ip, PyObject *op, PyArrayObject *out,
     if (multi == NULL) {
         goto fail;
     }
-    dtype = PyArray_DESCR(mps[0]);
-
-    /* Set-up return array */
-    if (out == NULL) {
-        Py_INCREF(dtype);
-        obj = (PyArrayObject *)PyArray_NewFromDescr(Py_TYPE(ap),
-                                                    dtype,
-                                                    multi->nd,
-                                                    multi->dimensions,
-                                                    NULL, NULL, 0,
-                                                    (PyObject *)ap);
-    }
-    else {
-        int flags = NPY_ARRAY_CARRAY |
-                    NPY_ARRAY_WRITEBACKIFCOPY |
-                    NPY_ARRAY_FORCECAST;
 
+    if (out != NULL) {
         if ((PyArray_NDIM(out) != multi->nd)
                     || !PyArray_CompareLists(PyArray_DIMS(out),
                                              multi->dimensions,
@@ -1044,24 +1029,18 @@ PyArray_Choose(PyArrayObject *ip, PyObject *op, PyArrayObject *out,
                             "choose: invalid shape for output array.");
             goto fail;
         }
+    }
 
-        for (i = 0; i < n; i++) {
-            if (arrays_overlap(out, mps[i])) {
-                flags |= NPY_ARRAY_ENSURECOPY;
-            }
-        }
+    dtype = PyArray_DESCR(mps[0]);
 
-        if (clipmode == NPY_RAISE) {
-            /*
-             * we need to make sure and get a copy
-             * so the input array is not changed
-             * before the error is called
-             */
-            flags |= NPY_ARRAY_ENSURECOPY;
-        }
-        Py_INCREF(dtype);
-        obj = (PyArrayObject *)PyArray_FromArray(out, dtype, flags);
-    }
+    /* Set-up return array */
+    Py_INCREF(dtype);
+    obj = (PyArrayObject *)PyArray_NewFromDescr(Py_TYPE(ap),
+                                                dtype,
+                                                multi->nd,
+                                                multi->dimensions,
+                                                NULL, NULL, 0,
+                                                (PyObject *)ap);
 
     if (obj == NULL) {
         goto fail;
@@ -1137,7 +1116,10 @@ PyArray_Choose(PyArrayObject *ip, PyObject *op, PyArrayObject *out,
     PyDataMem_FREE(mps);
     if (out != NULL && out != obj) {
         Py_INCREF(out);
-        PyArray_ResolveWritebackIfCopy(obj);
+        if (PyArray_CopyAnyInto(out, obj) < 0) {
+            Py_DECREF(out);
+            goto fail;
+        }
         Py_DECREF(obj);
         obj = out;
     }
diff --git a/numpy/_core/tests/test_multiarray.py b/numpy/_core/tests/test_multiarray.py
index 521cce26ea5a..098665fec5b6 100644
--- a/numpy/_core/tests/test_multiarray.py
+++ b/numpy/_core/tests/test_multiarray.py
@@ -10223,6 +10223,16 @@ def test_gh_24459():
         np.choose(a, [3, -1])
 
 
+def test_gh_28206():
+    a = np.arange(3)
+    b = np.ones((3, 3), dtype=np.int64)
+    out = np.array([np.nan, np.nan, np.nan])
+
+    with warnings.catch_warnings():
+        warnings.simplefilter("error", RuntimeWarning)
+        np.choose(a, b, out=out)
+
+
 @pytest.mark.parametrize("N", np.arange(2, 512))
 @pytest.mark.parametrize("dtype", [np.int16, np.uint16,
                         np.int32, np.uint32, np.int64, np.uint64])


Code review comments:
diff_hunk: @@ -1044,24 +1029,18 @@ PyArray_Choose(PyArrayObject *ip, PyObject *op, PyArrayObject *out,
                             "choose: invalid shape for output array.");
             goto fail;
         }
+    }
 
-        for (i = 0; i < n; i++) {
-            if (arrays_overlap(out, mps[i])) {
-                flags |= NPY_ARRAY_ENSURECOPY;
-            }
-        }
+    dtype = PyArray_DESCR(mps[0]);
 
-        if (clipmode == NPY_RAISE) {
-            /*
-             * we need to make sure and get a copy
-             * so the input array is not changed
-             * before the error is called
-             */
-            flags |= NPY_ARRAY_ENSURECOPY;
-        }
-        Py_INCREF(dtype);
-        obj = (PyArrayObject *)PyArray_FromArray(out, dtype, flags);
path: numpy/_core/src/multiarray/item_selection.c
text: @user1:
Hmmmm, agree that this fixes the issue.  There is somewhat large downside, though: If the user passed an array that does full-fill all the requirements (not using clipmode raise, C-contiguous and aligned, and identical dtype), then we will now always do an unnecessary copy.

And that is a bit unfortunate, because it means that passing `out=...` never saves anythign compared to just assigning manually.

@author:
Ah right, in the case where flags is 0, we saved the copy on PyArray_FromArray, and that optimization is missing now. That indeed somewhat removes a useful feature.

I'll revert the logic to the previous implementation and instead move around the conditions to include a check for the dtype. Thanks!

@author:
I modified the code to create a new array only when the dtype is different. Doesn't look too nice honestly, but is a commonsense way to optimize given our constraints.

@user2:
You could create a set of inputs to the function that doesn't create a copy on `main`, and in that case check to make sure that a copy doesn't happen. You can check the `base` attribute of an ndarray to see if it's a view or copy:

```
In [1]: import numpy as np

In [2]: arr = np.array([1, 2, 3])

In [3]: newarr = arr[:2]

In [4]: newarr
Out[4]: array([1, 2])

In [5]: newarr.base
Out[5]: array([1, 2, 3])

In [6]: newarr.base is arr
Out[6]: True
```

@user2:
And for a copy:

```
In [7]: newarr.copy()
Out[7]: array([1, 2])

In [8]: newarr.copy().base is None
True
```

@author:
Oh thank you for the guidance! That makes sense--on this.

@author:
Hmm. So on creating a test:

```python
def test_gh_28206():
    # no copy should be made in this case
    a = np.arange(3)
    b = np.ones((3, 3), dtype=np.int64)
    out = np.ones(3, dtype=np.int64)

    np.choose(a, b, out=out)
    assert out.base is not None
```

It fails on both the main branch and the branch for this PR. It seems that by using `PyArray_FromArray`, we may be using the same buffer, creating neither a view nor a copy? Should we check the array interface instead?

@user2:
Hmm, so maybe it's not so easy to test this:

```
In [1]: a = np.arange(3)
   ...: b = np.ones((3, 3), dtype=np.int64)

In [2]: out = np.empty_like(a)

In [3]: id(out)
Out[3]: 4363385904

In [4]: old_out = out

In [5]: ret = np.choose(a, b, out=out)

In [6]: id(out)
Out[6]: 4363385904

In [7]: old_out is out
Out[7]: True

In [8]: ret is out
Out[8]: True

In [9]: a = np.arange(3, dtype=np.int32)
   ...: b = np.ones((3, 3), dtype=np.int32)

In [10]: out = np.empty(3, dtype=np.int64)

In [11]: id(out)
Out[11]: 4364077264

In [12]: ret = np.choose(a, b, out=out)

In [13]: out
Out[13]: array([1, 1, 1])

In [14]: out.dtype
Out[14]: dtype('int64')

In [15]: id(out)
Out[15]: 4364077264

In [16]: out.base

In [17]: ret.dtype
Out[17]: dtype('int64')

In [18]: ret is out
Out[18]: True
```

Looking at the code, the return object is set to `out` and `out`'s descriptor gets replaced, so I guess this is all expected.

Never mind about the test.

@user1 do you see any issues with the new approach here after your last suggestion?

@user1:
This looks correct to me.  But now we can (and should) simplify it more.  We used the flags to enforce a copy when (please double check):
1. The dtype doesn't match
2. The array is not contiguous or not aligned
3. There is a memory overlap.

I really don't think point 2 matters.  You have covered the dtype now explicitly.  Which leaves the memory overlap, which we can just also take this branch!

@author:
Actually, it seems we did not check for C-contiguous! Instead, when we convert the inputs to a common type using `PyArray_ConvertToCommonType` that handles it :)

That leaves dtype and memory overlap. I'll take memory overlap!

@author:
Oh forgot about `clipmode=raise`. So, we have three conditions: dtype, memory overlap, and clipmode raise. Thinking about how to simplify that.

@author:
Okay, so it doesn't seem like there's a way to avoid creating a copy on `clipmode=raise`, so I guess what we can do is take both memory overlap and clipmode raise in this new branch.

@author:
Just pushed this!