BUG: Remove unnecessary copying and casting from out array in choose by MaanasArora · Pull Request #28206 · numpy/numpy

MaanasArora · 2025-01-21T10:49:18Z

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!

…_Choose

ngoldbaum · 2025-01-21T15:47:00Z

Neat! It's a good sign the tests pass. I don't know offhand if we have great test coverage for np.choose. One thing you could do is add a test to this PR based on the test code in the linked issue demonstrating that np.choose no longer generates a spurious warning. There's an example in the pytest docs showing how to convert warnings into errors so the test fails by raising an error if a warning is generated.

MaanasArora · 2025-01-21T16:27:45Z

@ngoldbaum Added the test, thanks! It fails in main and passes in this branch.

ngoldbaum

Just one small nitpick. I thought carefully about reference counts and didn't see any issues.

@seberg can you give this a once-over once the failure case for CopyAnyInto gets handled and hit the merge button assuming you don't see any additional issues?

ngoldbaum · 2025-01-22T22:30:13Z

numpy/_core/src/multiarray/item_selection.c

    if (out != NULL && out != obj) {
        Py_INCREF(out);
-        PyArray_ResolveWritebackIfCopy(obj);
+        PyArray_CopyAnyInto(out, obj);


This can fail so you need to check the failure case. The same bug existed in the code you're updating - ResolveWritebackIfCopy also returns -1 on failure.

I added a simple failure check; hopefully that is sufficient. Thank you!

ngoldbaum · 2025-01-23T16:49:34Z

numpy/_core/src/multiarray/item_selection.c

    if (out != NULL && out != obj) {
        Py_INCREF(out);
-        PyArray_ResolveWritebackIfCopy(obj);
+        if (PyArray_CopyAnyInto(out, obj) != 0) {


Suggested change

if (PyArray_CopyAnyInto(out, obj) != 0) {

if (PyArray_CopyAnyInto(out, obj) < 0) {

Slightly more defensive in case another non-error return code is added in the future (almost certainly won't happen but a good defensive programming habit to get into).

That makes sense; I made this change, thanks.

numpy/_core/src/multiarray/item_selection.c

…-unnecessary-cast

MaanasArora · 2025-01-31T07:55:20Z

Sorry, it seems some unrelated changes creeped in so I removed them. This new approach should be ready for review now. Thank you

ngoldbaum · 2025-01-31T15:08:29Z

Maybe also another test to make sure the unnecessary copy the old approach triggered doesn't happen?

MaanasArora · 2025-02-01T19:18:19Z

Not entirely sure if there's a convenient way to test for that in pytest? In either approach, there would probably be no data indicating a copy was made, after the array is overwritten.

MaanasArora · 2025-02-06T09:36:00Z

Sorry about the accidental changes--I'll be more careful with Git

seberg

Thanks a lot, yes I think this is correct now! I also think we can simplify things more, by just adding a must_copy flag and using your new branch then unconditionally.

seberg · 2025-02-06T10:21:33Z

numpy/_core/src/multiarray/item_selection.c

        Py_INCREF(dtype);
-        obj = (PyArrayObject *)PyArray_FromArray(out, dtype, flags);
+        if (!PyArray_EquivTypes(dtype, PyArray_DESCR(out))) {
+            obj = (PyArrayObject *)PyArray_NewFromDescr(Py_TYPE(out),


I don't think there is any reason to use Py_TYPE(out) here, just pass NULL and get a base-class array, I think.

(There is a subtle issue here about that, I think, in that ufuncs would ensure they call __array_wrap__ when done and we don't. But this a completely different issue!)

Makes sense thanks, but looking at the doc for PyArray_NewFromDescr, should we pass &PyArray_Type instead? NULL causes a ValueError.

Yes, of course. Somehow thought it accepted NULL as PyArray_Type there.

I made the change but it seems to fail a regression test test_attributes that tests for an identical info string assigned to the out array. Would you like me to push?

Maybe push, but may not look at it today. I really don't see why it matters since obj is never returned to the user! (And not understanding that worries me more than the change ;)).

No problem at all, thanks for the feedback. Yes, it is a bit confusing...

Oops, I think I was editing the wrong line! Things work fine now, pushed. Thanks.

numpy/_core/src/multiarray/item_selection.c

seberg · 2025-02-06T10:31:45Z

numpy/_core/src/multiarray/item_selection.c

-        PyArray_ResolveWritebackIfCopy(obj);
+        int success;
+        if (is_newarr) {
+            success = PyArray_CopyAnyInto(out, obj);


In theory, you could set the base-object to your allocation above and set to use WritebackIfCopy even with your custom allocation. But, TBH, I think a direct copy is actually clearer (frankly, I am not sure if we couldn't use it in half the places ResolveWritebackIfCopy is used ;)).

I see--good to know!

seberg · 2025-02-06T11:29:01Z

numpy/_core/tests/test_multiarray.py


+        # gh-28206 test new StringDType casting
+        x = np.array(['a', 'b', 'c'], dtype=np.dtypes.StringDType())
+        y = np.array(['abc', 'def', 'ghi'], dtype=np.dtypes.StringDType())


Subtle thing, but make sure to include a long string (otherwise you can't trigger the path which might have crashed).

Actually that still seems to crash, thanks for noticing. I think PyArray_ConvertToCommonType (it seems to be a backward compatibility function) does not cast properly to a common scenario for StringDType, and EquivTypes passes the dtype unchanged--so, we end up with the dtype only for the first array. Should we do a better promotion of some sort if the types are equivalent? Or perhaps, for now, we can ensure they are identical?

seberg · 2025-02-06T11:52:46Z

numpy/_core/src/multiarray/item_selection.c

+        }
+        if (success < 0) {
+            Py_DECREF(obj);
+            goto fail;


Just noticed: this is wrong! At this point the code already did most/all cleanup and doing that again will cause crashes. So you just need to do a straight return NULL.

Ah yes, the other functions do the cleanup already--fixed this!

seberg

Thanks, but one more iteration, I think. The idea here was to always manually copy out rather than use any auto-setup via PyArray_FromArray and that final writebackifcopy finalization.

So, I think both should be gone at the end, because otherwise it is confusing when what happenes.

seberg · 2025-02-07T18:59:26Z

numpy/_core/src/multiarray/item_selection.c

    }
    dtype = PyArray_DESCR(mps[0]);

+    int is_newarr = 0;


Can you find a different name for this? At the end of this, this should have the same meaning as out != NULL && obj != out. (or make it local and keep using that)

I.e. something clarifying that the user passed out array was copied (or must be copied).

I used copy_existing_out, which hopefully isn't too clumsy!

Wait, wasn't is_newarr maybe more fitting as it also applies when the array is created from scratch (in out == NULL)? It seems to be a special case of out != NULL && obj != out after out is (re)assigned. I'll do out_newarr maybe?

Maybe, in my mind it doesn't apply to out == NULL, because to me it tells us if we need to copy back at the end (as well).

Ah okay, I agree! What about copy_out? Since it's not always user passed...

seberg · 2025-02-07T19:01:30Z

numpy/_core/src/multiarray/item_selection.c

+                        NPY_ARRAY_WRITEBACKIFCOPY |
+                        NPY_ARRAY_FORCECAST;
+
+            obj = (PyArrayObject *)PyArray_FromArray(out, dtype, flags);


This branch should now never copy, so the only thing it should be doing is check for PyArray_FailUnlessWriteable. That also means you don't need the writeback-if-copy branch below.

(We can check that earlier also, since it doesn't make sense to wait until the last copy for this to fail in the other branch either.)

Makes sense, thanks! Addressed this.

seberg · 2025-02-07T19:03:00Z

numpy/_core/src/multiarray/item_selection.c

+        }
+        else {
+            success = PyArray_ResolveWritebackIfCopy(obj);
+        }


And thus, you never need this branch here either.

This as well.

seberg · 2025-02-10T11:33:06Z

numpy/_core/src/multiarray/item_selection.c

                                                    NULL, NULL, 0,
                                                    (PyObject *)ap);
-        out_newarr = 1;
+        copy_out = 1;


Suggested change

copy_out = 1;

Oh yes, your naming suggestion makes total sense now 😓 I'll change that back too, thanks!

seberg · 2025-02-10T11:35:13Z

numpy/_core/src/multiarray/item_selection.c

            goto fail;
        }

+        PyArray_FailUnlessWriteable(out, "output array");


Can you please add a test for this, because this doesn't work, you need an error return check and deal with that (this is not C++ with exceptions).

Ah yes sorry, fixed this, adding a test now.

Test is pushed.

seberg · 2025-02-10T11:38:01Z

numpy/_core/src/multiarray/item_selection.c

+                                                        (PyObject *)out);
+        }
+        else {
+            obj = (PyArrayObject *)PyArray_FromArray(out, dtype, 0);


Please only incref here and assign to obj. We rely on that being true now.

Sorry, just to be clear, this would be the dtype? Made that change.

seberg

LGTM, modulo some refactors (e.g. error refcounting was wrong).

numpy/_core/src/multiarray/item_selection.c

MaanasArora · 2025-02-25T09:20:09Z

Thank you for reviewing.

MaanasArora added 2 commits January 21, 2025 16:09

BUG: Remove unnecessary copying and casting from out array in PyArray…

ba7752a

…_Choose

BUG: Add missing output dimension check in PyArray_Choose

f7a9f15

github-actions bot added the 00 - Bug label Jan 21, 2025

Add test checking for spurious warning in np.choose

2df206d

ngoldbaum reviewed Jan 22, 2025

View reviewed changes

BUG: Add failure check to copying operation in np.choose

464733f

ngoldbaum reviewed Jan 23, 2025

View reviewed changes

BUG: Add defensive comparison instead of equals for failure check

dda7513

seberg reviewed Jan 24, 2025

View reviewed changes

numpy/_core/src/multiarray/item_selection.c Show resolved Hide resolved

MaanasArora added 7 commits January 25, 2025 10:24

BUG: Revert changes

ac537d1

BUG: Add check in np.choose for different dtypes to create new array

77da0b1

BUG: Revert accidental changes

239c2e9

BUG: Revert accidental changes

e62f259

Merge branch 'main' of https://github.com/numpy/numpy into bug/choose…

8681f85

…-unnecessary-cast

BUG: Revert accidental changes

7207a9d

BUG: Revert unnecessary changes

bec2ef4

MaanasArora added 2 commits February 6, 2025 14:49

ENH: Add test for large array fast powers

c32412f

BUG: Revert accidental changes

b7bbacb

seberg reviewed Feb 6, 2025

View reviewed changes

MaanasArora added 2 commits February 6, 2025 16:42

BUG: Adjust transfer function for cast in PyArray_Choose

71b5079

TST: Add test for string casting in np.choose

72e1709

seberg reviewed Feb 6, 2025

View reviewed changes

TST: Add long string to np.choose for StringDType

b3fb397

seberg reviewed Feb 6, 2025

View reviewed changes

MaanasArora added 3 commits February 6, 2025 17:26

BUG: Fix duplicate code cleanup in failure in np.choose

4ed5e06

TST: Remove np.choose string test to be addressed in future

841e232

MAINT: Simplify array copying logic in np.choose

27fed37

seberg reviewed Feb 7, 2025

View reviewed changes

MaanasArora added 4 commits February 8, 2025 16:44

MAINT: Simplify output array handling in np.choose

ea9a488

MAINT: Change confusing variable name

3bd5805

MAINT: Remove unnecessary type in new array from descr

fe45e50

MAINT: Rename confusing variable

d2d8bac

seberg reviewed Feb 10, 2025

View reviewed changes

MaanasArora added 4 commits February 10, 2025 17:15

BUG: Fix incorrect error handling

44c11e3

BUG: Remove unnecessary variable assignment and clarify name

0ef3a04

BUG: Remove reference increment not needed in conditional branch

5b09792

TST: Check fail when out in np.choose not writeable

8dde114

seberg approved these changes Feb 25, 2025

View reviewed changes

Apply suggestions from code review

b29ff77

seberg merged commit cdb44f7 into numpy:main Feb 26, 2025
61 of 62 checks passed

	if (PyArray_CopyAnyInto(out, obj) != 0) {
	if (PyArray_CopyAnyInto(out, obj) < 0) {

Uh oh!

Conversation

MaanasArora commented Jan 21, 2025

Uh oh!

ngoldbaum commented Jan 21, 2025

Uh oh!

MaanasArora commented Jan 21, 2025

Uh oh!

ngoldbaum left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

MaanasArora commented Jan 31, 2025

Uh oh!

ngoldbaum commented Jan 31, 2025

Uh oh!

MaanasArora commented Feb 1, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

MaanasArora commented Feb 6, 2025

Uh oh!

seberg left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

MaanasArora Feb 6, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

MaanasArora Feb 6, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

seberg left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

MaanasArora Feb 8, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

MaanasArora Feb 8, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

MaanasArora commented Feb 1, 2025 •

edited

Loading

MaanasArora Feb 6, 2025 •

edited

Loading

MaanasArora Feb 6, 2025 •

edited

Loading

MaanasArora Feb 8, 2025 •

edited

Loading

MaanasArora Feb 8, 2025 •

edited

Loading