NumPy 2.0 Support

[albanD]

NumPy 2.0 is expected to be out in December 2023. See details at numpy/numpy#24300

The major changes are going to be some Python API and behavior improvements. These don't impact PyTorch code itself and our compat layer (https://github.com/pytorch/pytorch/tree/main/torch/_numpy) is targeting 2.0+ already.

There are also C-API changes and in particular while 2.0 ABI will be backward compatible, the current ABI will not be forward compatible with 2.0. This will require some care on our end to make sure things work properly.

In terms of PyTorch future release, it means that, as of today, the expected support matrix looks like:

• PyTorch <=2.1 and nightly until EOY 2023 can live alongside any version of NumPy but the .numpy() and .from_numpy() functions will ONLY work if the version of NumPy installed is <2.0. Users can build from source against their version of NumPy if they need a more recent version of NumPy.
• PyTorch >=2.2 and nightly after EOY 2023 (once we have a NumPy 2.0 ABI-compatible build) will be fully compatible with all versions of NumPy.

The AIs to ensure this works:

• Add a helpful error message in PyTorch to give users a helpful error message if they try to use .numpy() or .from_numpy() with numpy 2.0+ today. This will replace the ABI error message from https://github.com/numpy/numpy/blob/7b56da8c46427d487b9337780417a9aad64eb295/numpy/core/code_generators/generate_numpy_api.py#L81-L82 that they will get today by a message asking them use NumPy <2.0 or build from source.
• Once we have a NumPy 2.0 compatible available, upgrade all our binary build to this version of NumPy and remove the error message above. This will ensure that our produced binaries work with any version of NumPy (including <2.0)
• Make torch.compile compatible with NumPy 2.0. In order to do that, we may need to update the test/torch_np tests that we pulled from NumPy

cc @mruberry @rgommers

[seemethere]

cc @atalman, @osalpekar

[malfet]

IMO we should do some of that for 2.3 (already started by making PyTorch compilable against Numpy-2.0)

[rgommers]

IMO we should do some of that for 2.3 (already started by making PyTorch compilable against Numpy-2.0)

I agree. It is important to build all binaries for the next minor or bugfix release of PyTorch against the latest NumPy 2.0.0 pre-release. That way, it will be ABI-compatible with both 1.x and 2.0, which is going to avoid a lot of problems. If such a release is not done before numpy 2.0.0 final is out (probably 8 weeks from now), then any environments created with pip install torch numpy will be quite broken, since no interop between PyTorch and NumPy is possible then. I haven't tried, but it's also possible that import torch or use of sub-components will be broken completely, since try-except imports of numpy that are scattered around the PyTorch codebase are going to succeed and then give RuntimeError('ABI mismatch ...') when used (example in TorchDynamo).

The plan posted by @albanD in the issue description still looks right. The first item is probably no longer relevant, since 2.0.0b1 is already available. The second item is urgent. The third one good to do, but the compat issues are probably relatively minor, so not nearly as urgent as handling the ABI compat issue.

[albanD]

Note that we still support 3.8+ in PyTorch. So only the version of CPython for which Numpy 2.0 will be released need to be updated. Based on the versions available for the beta, I guess we will need to change the 3.9+ builds of PyTorch only.

I think we should:

1. Locally ensure that all is good to compile and sanity checks
2. Update the trunk builds and make sure our CI is all good (either running old numpy or 2.0)
3. Cherry pick this into the release branch

@atalman I can take care of 1 today, would someone on the Release/Infra team be able to work on 2 and 3 so that we can get this done for the next RC?

[rgommers]

Note that numpy 2.0.0rc1 is out now, so the ABI is officially frozen. The only relevant different is that pybind11<2.12.0 raises an error, but PyTorch already upgraded to pybind11 2.12.0. So I think the build dependency can be upgraded from 2.0.0b1 (as done in pytorch/builder#1747) to 2.0.0rc1.

[albanD]

Thanks for the info Ralf!
@atalman what is the 2.3 status for this? Anything outstanding or we're on track?

[atalman]

@albanD : Binaries should be supporting rc1 for both wheels and conda packages now, since this PR is landed: pytorch/builder#1768 will be cherry-picking this into release and generating an rc

[albanD]

Tested the RC from https://dev-discuss.pytorch.org/t/pytorch-release-2-3-0-final-rc-is-available/1995 with both the 2.0 and 1.26 versions and everything is working smoothly.
Also version for python 3.8 is properly compatible with the 1.24 release of numpy that is the latest available there.
I'll keep this open until 2.3 is released successfully and then will close.

[jakirkham]

Thanks for all of your work here! 🙏

Binaries should be supporting rc1 for both wheels and conda packages now, since this PR is landed: pytorch/builder#1768 will be cherry-picking this into release and generating an rc

Would it be possible to move the NumPy Conda packages to a different label (like dev or similar)?

Currently users are picking them up unintentionally when installing from the pytorch channel: pytorch/builder#1776

Edit: To clarify the NumPy RC packages (not PyTorch's own packages) are what users are picking up

[atalman]

@jakirkham removed numpy from pytorch and pytorch-conda conda channel. Thank you for reporting this issue.
pytorch CI is currently using https://anaconda.org/atalman/numpy for numpy 2.0.0rc1 builds until anaconda releases numpy 2.0

[PaliC]

Validated pytorch 2.3.0 from https://download.pytorch.org/whl/test/cu121 works for numpy==2.0.0rc1 and numpy==1.26.4

[jakirkham]

removed numpy from pytorch and pytorch-conda conda channel. Thank you for reporting this issue. pytorch CI is currently using https://anaconda.org/atalman/numpy for numpy 2.0.0rc1 builds until anaconda releases numpy 2.0

Thanks @atalman! 🙏

Please feel free to close issue: pytorch/builder#1776

Also thanks for all your and the teams work on NumPy 2 support

Sounds good