ENH: support no-copy pickling for any array that can be transposed to a C-contiguous array #28105

IndifferentArea · 2025-01-06T04:51:11Z

As discussed in #26878, support no-copy pickling for any array that can be transposed to a C-contiguous array.

IndifferentArea · 2025-01-07T02:15:09Z

Hi, I’m new to contributing to NumPy and would appreciate some guidance on the following:

I’ve added a simple test for this feature. Are there additional test cases or scenarios I should consider?
Before requesting a review, should I document the changes in doc/release/upcoming_changes/ or elsewhere?

Thank you!

ngoldbaum · 2025-01-07T20:28:37Z

Before requesting a review, should I document the changes in doc/release/upcoming_changes/ or elsewhere?

Yes, see step 5 https://numpy.org/devdocs/dev/index.html#development-process-summary.

I'm also not totally sure what the policy is on changing the array pickle format, since people use pickles for uses they are not well-suited but we still need to support those uses anyway...

seberg

A few comments, we should discuss it but I think we can change it so long we make sure that the new NumPy can read the old pickles (and that also means adding a test that would currently fail, I think).

Also you need a test for the non-contiguous protocol=5 (unless there already is).

seberg · 2025-01-07T20:54:28Z

numpy/_core/numeric.py


-def _frombuffer(buf, dtype, shape, order):
-    return frombuffer(buf, dtype=dtype).reshape(shape, order=order)
+def _frombuffer(buf, dtype, shape, order, axis):


This function must still be able to deal with old pickles. It is effectively public API.

A test should have existed to find this. We should use pickle.dumps(...) on a current NumPy version, take that string and ensure that pickle.loads(...) on it still works (i.e. freeze the pickle in time).

This looks like it breaks existing pickles and we shouldn't.

I've given axis a default value None and added a test on a legacy pkl dumped in 2.2.1.

Should I keep this test in code? Or should I move it in test_regression.py?

seberg · 2025-01-07T20:55:26Z

numpy/_core/src/multiarray/methods.c

+    if (PyArray_IS_C_CONTIGUOUS((PyArrayObject *)self)) {
+        order = 'C';
+        picklebuf_args = Py_BuildValue("(O)", self);
+        rev_perm = Py_BuildValue("O", Py_None);


Might use PyTuple_Pack() where possible (even refactor).

I don't get it, you mean construct an empty tuple instead of None to rev_perm on c_contiguous and f_contiguous path?

I mean I would avoid Py_BuildValue when it is easy to do so, and it is here. But also OK to keep: this is inherited.

Applied as you suggested.

seberg · 2025-01-07T21:04:26Z

numpy/_core/src/multiarray/methods.c

+        for (int i = n - 1, check_s = 1; i >= 0; i--) {
+            if (check_s * PyArray_DESCR(self)->elsize !=
+                items[i].stride) {
+                return array_reduce_ex_regular(self, protocol);


If we assume it works usually and/or the other path is slow anyway, could also just not check this and see if the transpose works out?
Not sure that is better, though.

Yeah...Applied, this is much more straightforward!

Not sure if I Py_DECREF correctly. :(

numpy/_core/src/multiarray/methods.c

seberg · 2025-01-07T21:07:16Z

numpy/_core/src/multiarray/methods.c

+    } else {
+        shape = PyObject_GetAttrString((PyObject *)self, "shape");
+    }
+    return Py_BuildValue("N(NONNO)", from_buffer_func, buffer, (PyObject *)descr, shape,


This isn't right, you lose rev_perm.

what do you mean lose

leak? You need to decref it at some point or use N.

got it, thanks. never wrote on cpython api before, sorry for these oblivious problems.

IndifferentArea · 2025-01-09T05:56:04Z

Also you need a test for the non-contiguous protocol=5 (unless there already is).

There is already a test for that

numpy/numpy/_core/tests/test_multiarray.py

Lines 4372 to 4383 in 94132ac

    
           def test_non_contiguous_array(self): 
        
               non_contiguous_array = np.arange(12).reshape(3, 4)[:, :2] 
        
               assert not non_contiguous_array.flags.c_contiguous 
        
               assert not non_contiguous_array.flags.f_contiguous 
        
               # make sure non-contiguous arrays can be pickled-depickled 
        
               # using any protocol 
        
               for proto in range(2, pickle.HIGHEST_PROTOCOL + 1): 
        
                   depickled_non_contiguous_array = pickle.loads( 
        
                           pickle.dumps(non_contiguous_array, protocol=proto)) 
        
                   assert_equal(non_contiguous_array, depickled_non_contiguous_array)

seberg

Thanks, the new test looks good. A few more comment, and I have one more comment actually: I wonder if we should just pickle C/F contiguous ones without passing the axes order so that they will still unpickle on older versions (because it is easy, not because it is required).

seberg · 2025-01-09T08:39:24Z

numpy/_core/src/multiarray/methods.c

+    } else {
+        shape = PyObject_GetAttrString((PyObject *)self, "shape");
+    }
+    return Py_BuildValue("N(NONNO)", from_buffer_func, buffer, (PyObject *)descr, shape,


leak? You need to decref it at some point or use N.

seberg · 2025-01-09T08:49:45Z

numpy/_core/tests/test_multiarray.py

+        x = np.arange(24).reshape(3, 4, 2)
+        # generated by same code in 2.2.1
+        data_dir = os.path.join(os.path.dirname(__file__), 'data')
+        legacy_pkl_path = os.path.join(data_dir, "legacy_contiguous.pkl")


Thanks, this works (I suspect you could use dtype=np.int8 or so and just copy the whole bytes string into the code here also), but just if you like it.

I used byte str in code and added more examples

seberg · 2025-01-09T08:50:17Z

numpy/_core/src/multiarray/methods.c

+    if (PyArray_IS_C_CONTIGUOUS((PyArrayObject *)self)) {
+        order = 'C';
+        picklebuf_args = Py_BuildValue("(O)", self);
+        rev_perm = Py_BuildValue("O", Py_None);


I mean I would avoid Py_BuildValue when it is easy to do so, and it is here. But also OK to keep: this is inherited.

numpy/_core/src/multiarray/methods.c

seberg · 2025-01-09T08:55:46Z

numpy/_core/tests/test_multiarray.py

        assert_equal(f_contiguous_array, depickled_f_contiguous_array)
+
+    def test_transposed_contiguous_array(self):
+        transposed_contiguous_array = np.random.rand(2, 3, 4).transpose((1, 0, 2))


The transpose example is not chosen carefully, maybe you should just have multiple examples with pytest.parametrize and make at least one 4-5 dimensional so this is less likely to happen:

a = np.ones((1, 2, 3)) a.transpose(1, 0, 2).transpose(1, 0, 2).shape == a.shape

so I don't know from the test if the back-transpose is filled correctly.

added more test case

seberg · 2025-01-09T09:09:56Z

numpy/_core/src/multiarray/methods.c

+
+    PyObject* shape = NULL;
+    if(order == 'K') {
+        shape = PyObject_GetAttrString((PyObject *)transposed_array, "shape");


Old messy code, so you can ignore: but we have a PyArray_IntTupleFromIntp helper to use PyArray_NDIM and PyArray_DIMS directly.

thx, applied

numpy/_core/src/multiarray/methods.c

numpy/_core/numeric.py

numpy/_core/tests/test_multiarray.py

ngoldbaum

Spotted some issues. Please give this entire PR a careful once-over to make sure you're catching error cases and appropriately propagating errors when they happen. If you're feeling ambitious you could try writing test cases for the error cases, error cases tend to be poorly tested and buggy.

ngoldbaum · 2025-01-24T19:47:12Z

numpy/_core/tests/test_multiarray.py

+    @pytest.mark.parametrize('transposed_contiguous_array', 
+        [np.random.rand(2, 3, 4).transpose((1, 0, 2)), 
+         np.random.rand(2, 3, 4, 5).transpose((1, 3, 0, 2))] + 
+        [np.random.rand(*np.arange(2, 7)).transpose(np.random.permutation(5)) for _ in range(3)])


use a seeded RNG local to the test instead of the global np.random RNG

This is still using the global np.random RNG. If you want the test to be reproducible, you need to use a seeded RNG local to the test, e.g. using this function.

ngoldbaum · 2025-01-24T19:52:53Z

numpy/_core/src/multiarray/methods.c

+        if (transposed_array != NULL) {
+            Py_DECREF(transposed_array);
+        }


Suggested change

if (transposed_array != NULL) {

Py_DECREF(transposed_array);

}

Py_XDECREF(transposed_array);

This pattern is common enough it has its own macro in the C API. There are some other common operations that have an X variant that does NULL checking.

thx, applied

ngoldbaum · 2025-01-24T19:54:22Z

numpy/_core/src/multiarray/methods.c

+    else {
+        shape = PyArray_IntTupleFromIntp(PyArray_NDIM(self),
+                                         PyArray_SHAPE(self));
+    }


PyArray_IntTupleFromIntp can fail so you need to check if shape is NULL here.

ngoldbaum · 2025-01-24T19:55:51Z

numpy/_core/src/multiarray/methods.c

+        if (transposed_array != NULL) {
+            Py_DECREF(transposed_array);
+        }


Suggested change

if (transposed_array != NULL) {

Py_DECREF(transposed_array);

}

Py_XDECREF(transposed_array);

ngoldbaum · 2025-01-24T19:58:13Z

numpy/_core/src/multiarray/methods.c

-    picklebuf_class = PyObject_GetAttrString(pickle_module, "PickleBuffer");
-    Py_DECREF(pickle_module);
-    if (picklebuf_class == NULL) {
+    if (npy_cache_import_runtime("pickle", "PickleBuffer", &picklebuf_class)) {


Suggested change

if (npy_cache_import_runtime("pickle", "PickleBuffer", &picklebuf_class)) {

if (npy_cache_import_runtime("pickle", "PickleBuffer", &picklebuf_class) == -1) {

A little surprised the pickle tests pass given this bug...

Fixed. sorry i forgot checking the return value of npy_cache_import_runtime properly.

Since npy_cache_import_runtime will only return -1 or 0, I think it didn't actually cause a bug in the improper version so pytest didn't detect this.

numpy/_core/src/multiarray/methods.c

IndifferentArea · 2025-03-20T14:03:49Z

ping

ngoldbaum · 2025-03-24T16:16:53Z

Thanks for the ping and sorry for dropping this, I'm going to take another look at this.

ngoldbaum

Again, sorry for taking so long to get back to this.

This is looking much better than the last time I reviewed, thanks for the careful updates.

I'm approving but I want another experienced developer to look this over, since changing the pickle format has annoying backward compatibility concerns associated with making any mistakes.

@seberg or maybe @mhvk would either of you be interested in giving this another pass?

numpy/_core/src/multiarray/methods.c

seberg

Sorry for forgetting about this (it was great that you made many of the other smaller cleanups here!).

I think it's good, I have one substantial ask, unfortunately. Because I really think we should try to allow old NumPy versions to load most pickles and I don't think this does.

numpy/_core/src/multiarray/methods.c

IndifferentArea · 2025-03-26T14:07:47Z

seems something's wrong with freebsd's ci

seberg · 2025-03-26T14:09:02Z

@IndifferentArea don't worry about it. But if you would like to see it green, try merging/rebasing the main branch. That would likely fix it.

charris · 2025-05-14T18:34:21Z

Thanks @IndifferentArea

eendebakpt · 2025-05-14T19:25:08Z

@charris @seberg This PR introduced some linting errors on my local system. For example in numpy/_core/tests/test_multiarray.py line 4452 a trailing whitespace. The CI was green for the PR, in particular the lint test: https://github.com/numpy/numpy/actions/runs/14085614702/job/39448799495.

Not sure what happened here. The ruff version used in the CI is a bit older, but on my system it gives the same errors. I will create a PR to update ruff and address the trailing whitespace.

ngoldbaum · 2025-05-14T19:26:15Z

It probably needed a rebase before merging.

IndifferentArea · 2025-05-15T00:33:53Z

Awesome to see it merged!
been busy with school things recently :(
Thanks so much for your help with this pr, and I'm excited to contribute more in the future!

IndifferentArea added 2 commits January 6, 2025 12:38

add support for transposed contiguous array and corresponding test

0dcde49

Merge branch 'numpy:main' into numpyGH-26878

a27f5eb

github-actions bot added the 01 - Enhancement label Jan 6, 2025

IndifferentArea added 4 commits January 6, 2025 12:57

fix lint

50be134

rm bad includes

7cea83c

fix

bf5ecaa

fix

a566660

IndifferentArea changed the title ~~ENH: support no-copy pickling for any array that can be transposed to a C-contiguous array~~ WIP: ENH: support no-copy pickling for any array that can be transposed to a C-contiguous array Jan 6, 2025

IndifferentArea marked this pull request as draft January 6, 2025 05:54

discard vla

63da2b9

IndifferentArea changed the title ~~WIP: ENH: support no-copy pickling for any array that can be transposed to a C-contiguous array~~ ENH: support no-copy pickling for any array that can be transposed to a C-contiguous array Jan 7, 2025

IndifferentArea marked this pull request as ready for review January 7, 2025 02:15

seberg reviewed Jan 7, 2025

View reviewed changes

IndifferentArea added 3 commits January 9, 2025 13:36

accept legacy pkl

9494672

check return

6b9421a

fix ci

6ca72f1

IndifferentArea added 2 commits January 9, 2025 14:32

check after transpose

0f70349

discard auto include

c045bd9

seberg reviewed Jan 9, 2025

View reviewed changes

IndifferentArea added 8 commits January 10, 2025 13:05

use N format in pybuildvalue

aa357a9

use pytuple pack

aee895b

clean up

1b2f9d6

add comment on fall back behavior

4a16cd3

add mroe test

449ea59

use byte str instead of pkl to test

f0d528a

get shape with pyarray api

a8cc7fc

fmt

5f88eb3

ngoldbaum reviewed Jan 24, 2025

View reviewed changes

IndifferentArea added 8 commits January 27, 2025 18:08

use same random seed

bfbde0e

use xdecref

35ea29f

check PyArray_Transpose

596db8c

check PyArray_IntTupleFromIntp

1a82f4e

fix

80ba96d

once-over

792468a

add additional test

c213b37

once over again

3174d5e

IndifferentArea requested a review from ngoldbaum February 9, 2025 15:42

ngoldbaum approved these changes Mar 25, 2025

View reviewed changes

numpy/_core/src/multiarray/methods.c Outdated Show resolved Hide resolved

numpy/_core/src/multiarray/methods.c Outdated Show resolved Hide resolved

numpy/_core/src/multiarray/methods.c Outdated Show resolved Hide resolved

ngoldbaum added this to the 2.3.0 release milestone Mar 25, 2025

seberg reviewed Mar 25, 2025

View reviewed changes

numpy/_core/src/multiarray/methods.c Outdated Show resolved Hide resolved

numpy/_core/src/multiarray/methods.c Outdated Show resolved Hide resolved

numpy/_core/src/multiarray/methods.c Outdated Show resolved Hide resolved

IndifferentArea added 4 commits March 26, 2025 21:31

fmt

43e214a

pickle format consistency

42e1c99

fix incorrect usage of cpy api

f71549f

fmt

d5c2215

Merge branch 'main' into numpyGH-26878

8b981dc

charris merged commit d32cf93 into numpy:main May 14, 2025
72 checks passed

github-project-automation bot moved this from Awaiting a code review to Completed in NumPy first-time contributor PRs May 14, 2025

yangdong02 mentioned this pull request May 15, 2025

ENH: support no-copy pickling for any array that can be transposed to a C-contiguous array #26878

Closed

	if (npy_cache_import_runtime("pickle", "PickleBuffer", &picklebuf_class)) {
	if (npy_cache_import_runtime("pickle", "PickleBuffer", &picklebuf_class) == -1) {

Uh oh!

ENH: support no-copy pickling for any array that can be transposed to a C-contiguous array #28105

ENH: support no-copy pickling for any array that can be transposed to a C-contiguous array #28105

Uh oh!

Conversation

IndifferentArea commented Jan 6, 2025

Uh oh!

IndifferentArea commented Jan 7, 2025

Uh oh!

ngoldbaum commented Jan 7, 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!

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!

IndifferentArea Jan 9, 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!

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!

IndifferentArea Jan 10, 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!

IndifferentArea commented Jan 9, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

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!

Choose a reason for hiding this comment

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!

IndifferentArea Jan 10, 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!

Uh oh!

Uh oh!

Uh oh!

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

IndifferentArea Jan 9, 2025 •

edited

Loading

IndifferentArea Jan 10, 2025 •

edited

Loading

IndifferentArea commented Jan 9, 2025 •

edited

Loading

IndifferentArea Jan 10, 2025 •

edited

Loading

IndifferentArea Feb 2, 2025 •

edited

Loading