Make PyTensor compatible with numpy 2.0 #1194
Open
+943
−360
Conversation
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Codecov ReportAttention: Patch coverage is
❌ Your patch check has failed because the patch coverage (95.03%) 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.29% 82.32% +0.03%
==========================================
Files 186 187 +1
Lines 48034 48122 +88
Branches 8629 8639 +10
==========================================
+ Hits 39528 39615 +87
- Misses 6348 6350 +2
+ Partials 2158 2157 -1
|
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) |
|
PS: this is amazing 😍 I'll try to review soon |
|
We'll need to advertise slower advanced indexing operation on the C backend. I think that's the only regression? |
Yes, it 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. |
Also added ruff numpy2 transition rule.
Remaining tests now run on latest numpy, except for Numba jobs, which need numpy 2.1.0
- 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.
- 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
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
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. 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. The doctest for tensor.extra_ops.Unique was failing because the output shape for the inverse indices has changed when axis is None: numpy/numpy#20638
In numpy 2.0, if axis=None, then np.unique does not flatten the inverse indices returned if return_inverse=True The output shape has been changed to match numpy, dependent on the version of numpy.
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.
brendan-m-murphy
force-pushed
the
numpy-2-updates
branch
from
7ad30ee to
631f397
Compare
|
@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. |
13 tasks
|
@brendan-m-murphy Thanks for this pull request! Looking forward to it! |
ricardoV94
reviewed
Feb 11, 2025
Comment on lines
+162
to
+170
| - os: "ubuntu-latest" | ||
| python-version: "3.10" | ||
| numpy-version: "~=1.26.0" | ||
| fast-compile: 0 | ||
| float32: 0 | ||
| install-numba: 0 | ||
| install-jax: 0 | ||
| install-torch: 0 | ||
| part: "tests/tensor/test_math.py" |
There was a problem hiding this comment.
Why do we need this extra job?
| /* NumPy 2 definition always works, but add it for 1.x only */ | ||
| #define PyDataType_ISLEGACY(dtype) (1) | ||
| #else | ||
| #define NPY_DEFAULT_INT \\ |
There was a problem hiding this comment.
When can this final else be triggered? How is it possible to compile for two APIs at the same time?
| @@ -664,7 +763,7 @@ def c_init_code(self, **kwargs): | |||
| return ["import_array();"] | |||
|
|
|||
| def c_code_cache_version(self): | |||
| return (13, np.__version__) | |||
| return (18, np.version.git_revision) | |||
| @@ -3125,7 +3224,7 @@ def L_op(self, inputs, outputs, gout): | |||
| else: | |||
| return [x.zeros_like()] | |||
|
|
|||
| return (gz * exp2(x) * log(np.cast[x.type](2)),) | |||
| return (gz * exp2(x) * log(np.asarray(2, dtype=x.type)),) | |||
There was a problem hiding this comment.
Should all these be np.array(...) not np.asarray(...) since we know they are not arrays to start with?
Comment on lines
+19
to
+24
| try: | ||
| from numpy.lib.array_utils import normalize_axis_index, normalize_axis_tuple | ||
| except ModuleNotFoundError: | ||
| # numpy < 2.0 | ||
| from numpy.core.multiarray import normalize_axis_index | ||
| from numpy.core.numeric import normalize_axis_tuple |
There was a problem hiding this comment.
Can we put this in pytensor.utils and import from there to avoid defining it every time?
| outputs.append(typ()) | ||
| if NumpyVersion(np.__version__) >= "2.0.0rc1": | ||
| if axis is None: | ||
| inverse_shape = TensorType(dtype="int64", shape=x.type.shape) |
There was a problem hiding this comment.
nit: name this inverse_type not inverse_shape
| if self.axis is not None: | ||
| shape = (x_shape[axis],) | ||
| elif NumpyVersion(np.__version__) >= "2.0.0rc1": | ||
| shape = x_shape |
There was a problem hiding this comment.
Worried about the behavior changing depending on the numpy version. Can we keep the old 1.x behavior (which means we need to ravel this output in perform?). We can open an issue to add the 2.x behavior, which I think should be done with a sentinel flag that the user has to opt-in, at least until we deprecate numpy < 2.0 support
| @@ -101,7 +101,7 @@ def __init__( | |||
| if str(dtype) == "floatX": | |||
| self.dtype = config.floatX | |||
| else: | |||
| if np.obj2sctype(dtype) is None: | |||
| if np.dtype(dtype).type is None: | |||
There was a problem hiding this comment.
This needs a try/except now? np.dtype does not return None for invalid types, it raises
| assert max(n).dtype == dtype | ||
| i_max = eval_outputs(max(n)) | ||
| assert i_max == itype.max | ||
|
|
||
| @pytest.mark.xfail(reason="Fails due to #770") |
There was a problem hiding this comment.
you can have it combined with the parametrize above and skip just uint64 if you use pytest.params(...). You'll find other uses in the codebase
| assert i_max == data.max() | ||
|
|
||
| @pytest.mark.xfail(reason="Fails due to #770") | ||
| def test_uint64_special_value(self): |
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.
This suggestion is invalid because no changes were made to the code.
Suggestions cannot be applied while the pull request is closed.
Suggestions cannot be applied while viewing a subset of changes.
Only one suggestion per line can be applied in a batch.
Add this suggestion to a batch that can be applied as a single commit.
Applying suggestions on deleted lines is not supported.
You must change the existing code in this line in order to create a valid suggestion.
Outdated suggestions cannot be applied.
This suggestion has been applied or marked resolved.
Suggestions cannot be applied from pending reviews.
Suggestions cannot be applied on multi-line comments.
Suggestions cannot be applied while the pull request is queued to merge.
Suggestion cannot be applied right now. Please check back later.
Description
This PR makes PyTensor compatible with numpy versions >= 1.26 and < 2.2. (Numba requires numpy < 2.2.)
These changes include:
3will print, but in numpy >= 2.0,np.int64(3)will print)numpy.coreis nownumpy._core, and many functions have been moved fromcoreto new public locations likenumpy.lib; alsoAxisErrorneeds to be imported fromnumpy.exceptionsnow). These changes are conditional on numpy version number, except theAxisErrorimport, which is compatible with numpy 1.26.np.castis deprecated; its replacement isnp.asarray(..., dtype=...)np.uniquehas changed whenaxisisNone; this required changes in theUniqueOp. (This change is conditional on numpy version number.).astype. Conversions are no longer handled automatically; for instancenp.asarray(-1, dtype="unit8")will raise anOverflowError.TensorType.filteruses this new conversion method ifallow_downcastis true, which preserves the existing behaviorOverflowErrors (for numpy >= 2.0, orTypeErrorfor numpy < 2.0), or use equivalent but valid values (e.g. using255for auint8, instead of-1).NumpyAutocasterhas been changes to explicitly convert values to numpy types usingnp.asarray, which preserves the existing behavior. (The reason this preserves the behavior is that this is how the comparison is done inTensorType.filter, wherenp.asarray(data)is compared toconverted_data = np.asarray(data, self.dtype).)changed methods of
numpy.random.Generatorthat are used bycopyandpickle.Generatorwith independent state, you must usedeepcopynow, instead ofcopyGenerator.__getstate__()toNone. To get the state now, you must useGenerator.bit_generator.state.->elsizebyPyArray_ITEMSIZEScalarType. 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 typesnpy_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 ilkenpy_complex64to these types.np.MAXDIMSwas removed from the public API. This value was a common flag used to indicate thataxis=Nonehas been passed. Now there is an explicitly flagNPY_RAVEL_AXIS. Implementing this change was a bit variable across the affected code. A compatibility header was added topytensor/npy_2_compat.pyto makeNPY_RAVEL_AXISavailable for numpy < 2.0.MapIterwas removed from the public numpy C-API, so it was not possible to adapt the C-code forAdvancedIncSubtensor1; instead aNotImplementedErroris raised, so this Op defaults to the python implementation, which usesnp.add.at.npy_3k_compat.h): BUILD: clean out py2 stuff from npy_3kcompat.h numpy/numpy#26842Related Issue
Checklist
Type of change
📚 Documentation preview 📚: https://pytensor--1194.org.readthedocs.build/en/1194/