Make PyTensor compatible with numpy 2.0 #1194

brendan-m-murphy · 2025-02-05T14:42:18Z

Description

This PR makes PyTensor compatible with numpy versions >= 1.26 and < 2.2. (Numba requires numpy < 2.2.)

These changes include:

Unpinning numpy in environment files and pyproject
Adding a numpy 1.26 test job to CI
- runs on python 3.10, float32 0, fast-compile 0
- skips doctests, since it doesn't seem to easy to make these conditional on numpy version, and the numpy repr for numerical types has changed (e.g. in numpy < 2.0, 3 will print, but in numpy >= 2.0, np.int64(3) will print)
General numpy deprecations and namespace changes:
- updated imports due to changed name spaces (e.g. numpy.core is now numpy._core, and many functions have been moved from core to new public locations like numpy.lib; also AxisError needs to be imported from numpy.exceptions now). These changes are conditional on numpy version number, except the AxisError import, which is compatible with numpy 1.26.
- deprecations: the main change is that np.cast is deprecated; its replacement is np.asarray(..., dtype=...)
The return value of the inverse indices from np.unique has changed when axis is None; this required changes in the Unique Op. (This change is conditional on numpy version number.)
Changes in how overflows and type conversions are handled: to explicitly change the type of a numpy array, you must use .astype. Conversions are no longer handled automatically; for instance np.asarray(-1, dtype="unit8") will raise an OverflowError.
- TensorType.filter uses this new conversion method if allow_downcast is true, which preserves the existing behavior
- Several tests have been changes to either expect OverflowErrors (for numpy >= 2.0, or TypeError for numpy < 2.0), or use equivalent but valid values (e.g. using 255 for a uint8, instead of -1).
Changes to python type promotion (NEP 50)
- NEP 50 outlines changes in how python types are compared with and converted to numpy types. These changes were optional before numpy 2.0, but are now default. Essentially, the new rule is: if a python float is used in an operation with a numpy float, the type of the numpy float will always be used.
- The NumpyAutocaster has been changes to explicitly convert values to numpy types using np.asarray, which preserves the existing behavior. (The reason this preserves the behavior is that this is how the comparison is done in TensorType.filter, where np.asarray(data) is compared to converted_data = np.asarray(data, self.dtype).)
Changes to random number generators:
- The numpy PR numpy/numpy@44ba7ca
  changed methods of numpy.random.Generator that are used by copy and pickle.
- To get a copy of a numpy Generator with independent state, you must use deepcopy now, instead of copy
- This PR also changed the return value of Generator.__getstate__() to None. To get the state now, you must use Generator.bit_generator.state.
Changes due to changes in the numpy C-API
- Some minor changes with straightforward updates, e.g. replace ->elsize by PyArray_ITEMSIZE
- Changes to complex scalars in ScalarType. The numpy implementation of complex numbers has been changed from a struct with real and imaginary values to the native C-99 complex types. On disk, these are equivalent, but the real and imaginary parts C-99 complex types cannot be accessed using pointers. Numpy provides some macros to make accessing real and imaginary parts uniform across pre and post 2.0 version. Since these are implemented in terms of the types npy_cfloat, npy_cdouble, npy_clongdouble, some generic functions were added to the C code so that we do not need to explicitly translation bit size aliases ilke npy_complex64 to these types.
- The constant np.MAXDIMS was removed from the public API. This value was a common flag used to indicate that axis=None has been passed. Now there is an explicitly flag NPY_RAVEL_AXIS. Implementing this change was a bit variable across the affected code. A compatibility header was added to pytensor/npy_2_compat.py to make NPY_RAVEL_AXIS available for numpy < 2.0.
- MapIter was removed from the public numpy C-API, so it was not possible to adapt the C-code for AdvancedIncSubtensor1; instead a NotImplementedError is raised, so this Op defaults to the python implementation, which uses np.add.at.
Dropped support for Python 2 in C code.
- Numpy is starting to remove Python 2 compatibility code (from npy_3k_compat.h): BUILD: clean out py2 stuff from npy_3kcompat.h numpy/numpy#26842
- I have updated the lazylinker C code to use Python 3 exclusively.

Related Issue

Closes Test on numpy 2.0 #688
Related to Numpy 2.0 #689 (Draft PR with discussion)

Checklist

Checked that the pre-commit linting/style checks pass
Included tests that prove the fix is effective or that the new feature works
Added necessary documentation (docstrings and/or example notebooks)
If you are a pro: each commit corresponds to a relevant logical change

Type of change

📚 Documentation preview 📚: https://pytensor--1194.org.readthedocs.build/en/1194/

codecov · 2025-02-05T15:21:32Z

Codecov Report

Attention: Patch coverage is 96.77419% with 5 lines in your changes missing coverage. Please review.

Project coverage is 82.04%. Comparing base (bbe663d) to head (b633bca).
Report is 19 commits behind head on main.

Files with missing lines	Patch %	Lines
pytensor/sparse/rewriting.py	75.00%	2 Missing ⚠️
pytensor/link/numba/dispatch/random.py	50.00%	1 Missing ⚠️
pytensor/tensor/basic.py	75.00%	1 Missing ⚠️
pytensor/tensor/type.py	90.00%	1 Missing ⚠️

❌ Your patch status has failed because the patch coverage (96.77%) is below the target coverage (100.00%). You can increase the patch coverage or adjust the target coverage.

Additional details and impacted files

@@            Coverage Diff             @@
##             main    #1194      +/-   ##
==========================================
+ Coverage   82.01%   82.04%   +0.03%     
==========================================
  Files         187      188       +1     
  Lines       48467    48533      +66     
  Branches     8669     8675       +6     
==========================================
+ Hits        39749    39819      +70     
+ Misses       6554     6553       -1     
+ Partials     2164     2161       -3

Files with missing lines	Coverage Δ
pytensor/link/c/basic.py	`87.72% <100.00%> (+0.24%)`	⬆️
pytensor/link/c/interface.py	`90.09% <100.00%> (ø)`
pytensor/link/c/lazylinker_c.py	`75.30% <100.00%> (ø)`
pytensor/link/jax/dispatch/random.py	`93.70% <100.00%> (ø)`
pytensor/link/numba/dispatch/elemwise.py	`94.67% <100.00%> (ø)`
pytensor/npy_2_compat.py	`100.00% <100.00%> (ø)`
pytensor/scalar/basic.py	`80.68% <100.00%> (+0.09%)`	⬆️
pytensor/sparse/basic.py	`82.57% <100.00%> (ø)`
pytensor/tensor/blas.py	`73.95% <100.00%> (ø)`
pytensor/tensor/blas_headers.py	`70.78% <100.00%> (ø)`
... and 16 more

ricardoV94 · 2025-02-05T16:06:54Z

Numba requires numpy < 2.2.

Numba is an optional dependency for us so no need to pin, we can just make sure the CI runs with that version for the numba + benchmark tests (they are already separate)

ricardoV94 · 2025-02-05T16:15:38Z

PS: this is amazing 😍 I'll try to review soon

ricardoV94 · 2025-02-05T19:49:01Z

We'll need to advertise slower advanced indexing operation on the C backend. I think that's the only regression?

brendan-m-murphy · 2025-02-06T10:00:57Z

We'll need to advertise slower advanced indexing operation on the C backend. I think that's the only regression?

Yes, it np.add.at is possibly slower in some cases: aesara-devs/aesara#1506. There is an open issue tracking enhancements: numpy/numpy#23176

I haven't looked at any benchmarks for this change, so I'm not sure what the impact is.

I'm making some changes to unpin numpy and make sure the CI runs with 2.1 for numba, I'll push those today.

brendan-m-murphy · 2025-02-06T15:14:48Z

@ricardoV94 Okay, unpinned numpy and changed the CI to match. I ran the tests on my fork and they passed. There was some random mypy issue, so I've fixed that. Not sure why it was raised now.

metrizable · 2025-02-10T23:55:06Z

@brendan-m-murphy Thanks for this pull request! Looking forward to it!

ricardoV94

This looks great, some small comments / questions

.github/workflows/test.yml

pytensor/npy_2_compat.py

pytensor/scalar/basic.py

pytensor/tensor/basic.py

pytensor/tensor/extra_ops.py

pytensor/tensor/type.py

tests/tensor/test_math.py

brendan-m-murphy · 2025-02-14T14:22:43Z

This looks great, some small comments / questions

Okay, I think I've addressed everything. (Sorry I didn't realise when I was pushing to my fork that it was pushing here too... but the latest push should pass CI.)

I've moved most of the numpy-version conditional code to npy_2_compat.py, I changed Unique to have the same behavior as in numpy < 2.0, and I made the other changes we discussed above.

maresb · 2025-02-14T16:45:32Z

@brendan-m-murphy, just in case you didn't notice there's a merge conflict.

brendan-m-murphy · 2025-02-14T16:54:23Z

@brendan-m-murphy, just in case you didn't notice there's a merge conflict.

Thanks, I rebased onto main, should be good now.

maresb · 2025-02-14T17:31:32Z

@brendan-m-murphy, looks like there's still an unresolved merge conflict in your copy of pyproject.toml:

pytensor/pyproject.toml

Lines 131 to 136 in ecfd045

    
           [tool.ruff.lint] 
        
           <<<<<<< HEAD 
        
           select = ["B905", "C", "E", "F", "I", "UP", "W", "RUF", "PERF", "PTH", "ISC", "T20"] 
        
           ======= 
        
           select = ["B905", "C", "E", "F", "I", "UP", "W", "RUF", "PERF", "PTH", "ISC", "NPY201"] 
        
           >>>>>>> ca2dd3dce (Unpinned numpy)

ricardoV94 · 2025-02-14T18:11:16Z

@brendan-m-murphy It seems you resolved all the comments, so we're pretty much ready to merge once those last conflicts get solved!.

I will give another pass next week to see if we didn't miss anything else.

Repeating myself, but this is an incredible contribution 🙏

tests/tensor/test_io.py

ricardoV94 · 2025-02-17T14:53:20Z

Good news is all tests are passing!

jessegrabowski · 2025-02-17T14:54:23Z

sorry team :(

- replaced np.AxisError with np.exceptions.AxisError - the `numpy.core` submodule has been renamed to `numpy._core` - some parts of `numpy.core` have been moved to `numpy.lib.array_utils` Except for `AxisError`, the updated imports are conditional on the version of numpy, so the imports should work for numpy >= 1.26. The conditional imports have been added to `npy_2_compat.py`, so the imports elsewhere are unconditonal.

- Replace np.cast with np.asarray: in numpy 2.0, `np.cast[new_dtype](arr)` is deprecated. The literal replacement is `np.asarray(arr, dtype=new_dtype)`. - Replace np.sctype2char and np.obj2sctype. Added try/except to handle change in behavior of `np.dtype` - Replace np.find_common_type with np.result_type Further changes to `TensorType`: TensorType.dtype must be a string, so the code has been changed from `self.dtype = np.dtype(dtype).type`, where the right-hand side is of type `np.generic`, to `self.dtype = str(np.dtype(dtype))`, where the right-hand side is a string that satisfies: `self.dtype == str(np.dtype(self.dtype))` This doesn't change the behavior of `np.array(..., dtype=self.dtype)` etc.

Some macros were removed from npy_3k_compat.h. Following numpy, I updated the affected functions to the Python 3 names, and removed support for Python 2. Also updated lazylinker_c version to indicate substantial changes to the C code.

- replace `->elsize` by `PyArray_ITEMSIZE` - don't use deprecated PyArray_MoveInto

Anything `Hashable` should work, but I've made the return type `tuple[Hashable]` to keep with the current style. This means, e.g., we can use strings in the cache version.

This is done using C++ generic functions to get/set the real/imag parts of complex numbers. This gives us an easy way to support Numpy v < 2.0, and allows the type underlying the bit width types, like pytensor_complex128, to be correctly inferred from the numpy complex types they inherit from. Updated pytensor_complex struct to use get/set real/imag aliases defined above. Also updated operators such as `Abs` to use get_real, get_imag. Macros have been added to ensure compatibility with numpy < 2.0 Note: redefining the complex arithmetic here means that we aren't treating NaNs and infinities as carefully as the C99 standard suggets (see Appendix G of the standard). The code has been like this since it was added to Theano, so we're keeping the existing behavior.

MapIter was removed from the public numpy C-API in version 2.0, so we raise a not implemented error to default to the python code for the AdvancedInSubtensor1. The python version, defined in `AdvancedInSubtensor1.perform` calls `np.add.at`, which uses `MapIter` behind the scenes. There is active development on Numpy to improve the efficiency of `np.add.at`. To skip the C implementation and use the Python implementation, we raise a NotImplementedError for this op's c code if numpy>=2.0.

This was done for the python linker and numba linker. deepcopy seems to be the recommended method for copying a numpy Generator. After this numpy PR: numpy/numpy@44ba7ca `copy` didn't seem to actually make an independent copy of the `np.random.Generator` objects spawned by `RandomStream`. This was causing the "test values" computed by e.g. `RandomStream.uniform` to increment the RNG state, which was causing tests that rely on `RandomStream` to fail. Here is some related discussion: numpy/numpy#24086 I didn't see any official documentation about a change in numpy that would make copy stop working.

numpy.random.Generator.__getstate__() now returns none; to see the state of the bit generator, you need to use Generator.bit_generator.state. This change affects `RandomGeneratorType`, and several of the random tests (including some for Jax.)

`np.MAXDIMS` was removed from the public API and no replacement is given in the migration docs. In numpy <= 1.26, the value of `np.MAXDIMS` was 32. This was often used as a flag to mean `axis=None`. In numpy >= 2.0, the maximum number of dims of an array has been increased to 64; simultaneously, a constant `NPY_RAVEL_AXIS` was added to the C-API to indicate that `axis=None`. In most cases, the use of `np.MAXDIMS` to check for `axis=None` can be replaced by the new constant `NPY_RAVEL_AXIS`. To make this constant accessible when using numpy <= 1.26, I added a function to insert `npy_2_compat.h` into the support code for the affected ops.

In numpy 2.0, -1 as uint8 is out of bounds, whereas previously it would be converted to 255. This affected the test helper function `reduced_bitwise_and`. The helper function was changed to use 255 instead of -1 if the dtype was uint8, since this is what is needed to match the behavior of the "bitwise and" op. `reduced_bitwise_and` was only used by `TestCAReduce` in `tests/tensor/test_elemwise.py`, so it was moved there from `tests/tensor/test_math.py`

1. Changed autocaster due to new promotion rules With "weak promotion" of python types in Numpy 2.0, the statement `1.1 == np.asarray(1.1).astype('float32')` is True, whereas in Numpy 1.26, it was false. However, in numpy 1.26, `1.1 == np.asarray([1.1]).astype('float32')` was true, so the scalar behavior and array behavior are the same in Numpy 2.0, while they were different in numpy 1.26. Essentially, in Numpy 2.0, if python floats are used in operations with numpy floats or arrays, then the type of the numpy object will be used (i.e. the python value will be treated as the type of the numpy objects). To preserve the behavior of `NumpyAutocaster` from numpy <= 1.26, I've added an explicit conversion of the value to be converted to a numpy type using `np.asarray` during the check that decides what dtype to cast to. 2. Updates due to new numpy conversion rules for out-of-bounds python ints In numpy 2.0, out of bounds python ints will not be automatically converted, and will raise an `OverflowError` instead. For instance, converting 255 to int8 will raise an error, instead of returning -1. To explicitly force conversion, we must use `np.asarray(value).astype(dtype)`, rather than `np.asarray(value, dtype=dtype)`. The code in `TensorType.filter` has been changed to the new recommended way to downcast, and the error type caught by some tests has been changed to OverflowError from TypeError

I was getting a NameError from the list comprehensions saying that e.g. `pytensor_scalar` was not defined. I'm not sure why, but this is another (more verbose) way to do the same thing.

From numpy PR numpy/numpy#22449, the repr of scalar values has changed, e.g. from "1" to "np.int64(1)", which caused two doctests to fail.

In numpy 2.0, if axis=None, then np.unique does not flatten the inverse indices returned if return_inverse=True A helper function has been added to npy_2_compat.py to mimic the output of `np.unique` from version of numpy before 2.0

Due to changes in numpy conversion rules (NEP 50), overflows are not ignored; in particular, negating a unsigned int causes an overflow error. The test for `neg` has been changed to check that this error is raised.

I split this test up to test uint64 separately, since this is the case discussed in Issue pymc-devs#770. I also added a test for the exact example used in that issue. The uint dtypes with lower precision should pass. The uint64 case started passing for me locally on Mac OSX, but still fails on CI. I'm not sure why this is, but at least the test will be more specific now if it fails in the future.

Also added ruff numpy2 transition rule.

Remaining tests now run on latest numpy, except for Numba jobs, which need numpy 2.1.0

brendan-m-murphy · 2025-02-17T15:37:56Z

I rebased onto main, it doesn't look like there were any serious conflicts :)

maresb · 2025-02-17T16:27:13Z

All green except for codecov. Quick, let's get this in! 😀

ricardoV94 · 2025-02-17T16:32:54Z

codecov is never green

twiecki · 2025-02-19T15:52:11Z

This is a meaty one - congrats!

brendan-m-murphy mentioned this pull request Feb 5, 2025

Numpy 2.0 #689

Closed

11 tasks

ricardoV94 added major NumPy compatibility labels Feb 5, 2025

brendan-m-murphy force-pushed the numpy-2-updates branch from 7ad30ee to 631f397 Compare February 6, 2025 11:05

Mathnerd314 mentioned this pull request Feb 9, 2025

python312Packages.pymc: 5.19.1 -> 5.20.0 NixOS/nixpkgs#371939

Draft

13 tasks

maresb mentioned this pull request Feb 9, 2025

Support Python 3.13 #1199

Merged

11 tasks

ricardoV94 reviewed Feb 11, 2025

View reviewed changes

ricardoV94 changed the title ~~Numpy 2.0 updates~~ Make PyTensor compatible with numpy 2.0 Feb 11, 2025

brendan-m-murphy force-pushed the numpy-2-updates branch from 8c17e93 to 1f80278 Compare February 14, 2025 10:33

prusnak mentioned this pull request Feb 14, 2025

fix release-cuda eval errors NixOS/nixpkgs#379768

Open

13 tasks

brendan-m-murphy force-pushed the numpy-2-updates branch 3 times, most recently from 340f24c to e7b728f Compare February 14, 2025 14:11

brendan-m-murphy force-pushed the numpy-2-updates branch from e7b728f to ecfd045 Compare February 14, 2025 16:53

ricardoV94 mentioned this pull request Feb 14, 2025

Remove redundant code after numpy < 2.0 is deprecated #1211

Open

brendan-m-murphy force-pushed the numpy-2-updates branch from ecfd045 to b081c3a Compare February 17, 2025 11:37

ricardoV94 mentioned this pull request Feb 17, 2025

Deprecate batched blas stuff #1215

Open

ricardoV94 reviewed Feb 17, 2025

View reviewed changes

tests/tensor/test_io.py Show resolved Hide resolved

ricardoV94 and others added 19 commits February 17, 2025 15:31

Updated lazylinker C code

910b27c

Some macros were removed from npy_3k_compat.h. Following numpy, I updated the affected functions to the Python 3 names, and removed support for Python 2. Also updated lazylinker_c version to indicate substantial changes to the C code.

Changes for deprecations in numpy 2.0 C-API

92d96ff

- replace `->elsize` by `PyArray_ITEMSIZE` - don't use deprecated PyArray_MoveInto

Update type hint for c_code_cache_version

b20f401

Anything `Hashable` should work, but I've made the return type `tuple[Hashable]` to keep with the current style. This means, e.g., we can use strings in the cache version.

Fix for NameError in test

bce3613

I was getting a NameError from the list comprehensions saying that e.g. `pytensor_scalar` was not defined. I'm not sure why, but this is another (more verbose) way to do the same thing.

Updated doctests

45c3a01

From numpy PR numpy/numpy#22449, the repr of scalar values has changed, e.g. from "1" to "np.int64(1)", which caused two doctests to fail.

Preserve numpy < 2.0 Unique inverse output shape

2bfe6dd

In numpy 2.0, if axis=None, then np.unique does not flatten the inverse indices returned if return_inverse=True A helper function has been added to npy_2_compat.py to mimic the output of `np.unique` from version of numpy before 2.0

Fix test for neg on unsigned

cd75f95

Due to changes in numpy conversion rules (NEP 50), overflows are not ignored; in particular, negating a unsigned int causes an overflow error. The test for `neg` has been changed to check that this error is raised.

Unpinned numpy

720568c

Also added ruff numpy2 transition rule.

Added numpy 1.26.* to CI

b633bca

Remaining tests now run on latest numpy, except for Numba jobs, which need numpy 2.1.0

brendan-m-murphy force-pushed the numpy-2-updates branch from 4dcaab3 to b633bca Compare February 17, 2025 15:36

ricardoV94 merged commit 51ea1a0 into pymc-devs:main Feb 17, 2025
72 of 73 checks passed

maresb mentioned this pull request Feb 17, 2025

Drop numpy < 2 #1216

Open

brendan-m-murphy mentioned this pull request Feb 18, 2025

Make OpenGHG compatible with numpy 2.0 openghg/openghg#1042

Open

metrizable mentioned this pull request Feb 19, 2025

BUG: Not working with numpy 2+ (np.obj2sctype was removed in the NumPy 2.0 release.) pymc-devs/pymc#7592

Open

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Make PyTensor compatible with numpy 2.0 #1194

Make PyTensor compatible with numpy 2.0 #1194

brendan-m-murphy commented Feb 5, 2025 •

edited by github-actions bot

Loading

codecov bot commented Feb 5, 2025 •

edited

Loading

ricardoV94 commented Feb 5, 2025

ricardoV94 commented Feb 5, 2025

ricardoV94 commented Feb 5, 2025 •

edited

Loading

brendan-m-murphy commented Feb 6, 2025

brendan-m-murphy commented Feb 6, 2025 •

edited

Loading

metrizable commented Feb 10, 2025

ricardoV94 left a comment

brendan-m-murphy commented Feb 14, 2025

maresb commented Feb 14, 2025

brendan-m-murphy commented Feb 14, 2025

maresb commented Feb 14, 2025 •

edited

Loading

ricardoV94 commented Feb 14, 2025 •

edited

Loading

ricardoV94 commented Feb 17, 2025

jessegrabowski commented Feb 17, 2025

brendan-m-murphy commented Feb 17, 2025

maresb commented Feb 17, 2025

ricardoV94 commented Feb 17, 2025

twiecki commented Feb 19, 2025

Make PyTensor compatible with numpy 2.0 #1194

Make PyTensor compatible with numpy 2.0 #1194

Conversation

brendan-m-murphy commented Feb 5, 2025 • edited by github-actions bot Loading

Description

Related Issue

Checklist

Type of change

codecov bot commented Feb 5, 2025 • edited Loading

Codecov Report

ricardoV94 commented Feb 5, 2025

ricardoV94 commented Feb 5, 2025

ricardoV94 commented Feb 5, 2025 • edited Loading

brendan-m-murphy commented Feb 6, 2025

brendan-m-murphy commented Feb 6, 2025 • edited Loading

metrizable commented Feb 10, 2025

ricardoV94 left a comment

Choose a reason for hiding this comment

brendan-m-murphy commented Feb 14, 2025

maresb commented Feb 14, 2025

brendan-m-murphy commented Feb 14, 2025

maresb commented Feb 14, 2025 • edited Loading

ricardoV94 commented Feb 14, 2025 • edited Loading

ricardoV94 commented Feb 17, 2025

jessegrabowski commented Feb 17, 2025

brendan-m-murphy commented Feb 17, 2025

maresb commented Feb 17, 2025

ricardoV94 commented Feb 17, 2025

twiecki commented Feb 19, 2025

brendan-m-murphy commented Feb 5, 2025 •

edited by github-actions bot

Loading

codecov bot commented Feb 5, 2025 •

edited

Loading

ricardoV94 commented Feb 5, 2025 •

edited

Loading

brendan-m-murphy commented Feb 6, 2025 •

edited

Loading

maresb commented Feb 14, 2025 •

edited

Loading

ricardoV94 commented Feb 14, 2025 •

edited

Loading