Numpy 2.0 compatibility

[rflamary]

🚀 Feature

We need to upadte the CI for wheels building so that they compile wheel for numpy <2.0 and numpy >=2.0. Maybe we need a versioning scheme that allows those versions (with the numpy2.0 the most up to date but enforce the numpy version?)

Motivation

Pitch

Alternatives

Additional context

[matthewfeickert]

As NumPy v2.0 is out now, this would be a great boon to have. Are there particular technical issues with getting a new build out that need resolving? Or is this more a matter of (understandably) limited developer time at the moment?

[rflamary]

It is more a question of limited time ;). Actually I would expect POT to work on numpy 2.0 if build from source.

The idea is toccheck that all tests pass on numpy 2.0 and update the build wheel github action to run with two different versions of numpy. Then I guess we can do a proper release

[matthewfeickert]

The idea is to check that all tests pass on numpy 2.0

@rflamary This becomes an understandably difficult problem though, as in the tests you are also testing in an environment with torch, jax, and tensorflow.

POT/requirements.txt

Lines 8 to 11 in cba9c7b

	torch
	jax<=0.4.24
	jaxlib<=0.4.24
	tensorflow

The PyTorch and JAX (via ml-dtypes) teams already got NumPy 2.0 compatible wheels out, but as of 2024-06-18 TensorFlow does not. As POT tests against all these libraries, does that mean that POT doesn't plan to support NumPy 2.0 untl TensorFlow does (which might be weeks to months later)? If that is the plan, is it possible to run on a subset of your tests with NumPy 2.0 compatible dependencies to determine if things are functioning as expected, and then release wheels built to support NumPy?

[matthewfeickert]

Actually I would expect POT to work on numpy 2.0 if build from source.

Unfortunately, not from the v0.9.3 sdist:

# docker run --rm -ti python:3.12 /bin/bash
$ python -m pip --quiet install --upgrade uv
$ uv venv && . .venv/bin/activate
$ uv pip install --upgrade 'numpy>=2.0.0'
$ uv pip install --no-binary pot pot
Resolved 3 packages in 582ms
   Built pot==0.9.3
Downloaded 2 packages in 8.31s
Installed 2 packages in 18ms
 + pot==0.9.3
 + scipy==1.13.1
$ uv pip list
Package Version
------- -------
numpy   2.0.0
pot     0.9.3
scipy   1.13.1
$ python -c 'import ot'

A module that was compiled using NumPy 1.x cannot be run in
NumPy 2.0.0 as it may crash. To support both 1.x and 2.x
versions of NumPy, modules must be compiled with NumPy 2.0.
Some module may need to rebuild instead e.g. with 'pybind11>=2.12'.

If you are a user of the module, the easiest solution will be to
downgrade to 'numpy<2' or try to upgrade the affected module.
We expect that some modules will need time to support NumPy 2.

Traceback (most recent call last):  File "<string>", line 1, in <module>
  File "/.venv/lib/python3.12/site-packages/ot/__init__.py", line 21, in <module>
    from . import lp
  File "/.venv/lib/python3.12/site-packages/ot/lp/__init__.py", line 23, in <module>
    from .emd_wrap import emd_c, check_result, emd_1d_sorted
Traceback (most recent call last):
  File "<string>", line 1, in <module>
  File "/.venv/lib/python3.12/site-packages/ot/__init__.py", line 21, in <module>
    from . import lp
  File "/.venv/lib/python3.12/site-packages/ot/lp/__init__.py", line 23, in <module>
    from .emd_wrap import emd_c, check_result, emd_1d_sorted
  File "ot/lp/emd_wrap.pyx", line 1, in init ot.lp.emd_wrap
ImportError: numpy.core.multiarray failed to import (auto-generated because you didn't call 'numpy.import_array()' after cimporting numpy; use '<void>numpy._import_array' to disable if you are certain you don't need it).

[rflamary]

damnit you are right. No worry we will build wheel independently of jax or tensorflow (they depend only on numpy) but indeed we will indedd need to rework the tests also....

[rflamary]

hum i think this is because of the use of oldest_suported_numpy that build using older numpy:

POT/setup.py

Line 72 in cba9c7b

setup_requires=["oldest-supported-numpy", "cython>=0.23"],

[matthewfeickert]

Yeah, to be clear I am no expert here, but from quickly reading bits of the "For downstream package authors" NumPy guide (the NumPy 2.0-specific advice), I think oldest-supported-numpy in

POT/pyproject.toml

Line 2 in cba9c7b

requires = ["setuptools", "wheel", "oldest-supported-numpy", "cython>=0.23"]

needs to get swapped for "numpy>=2.0.0" to build wheels compatible with both NumPy 1.xx and 2.x. Though I haven't tested that.

[rflamary]

OK we will have a look but feel free to do a PR in the meantime (it should be changed in the toml and setup.py files I think). This might be an easier fix than expected

[matthewfeickert]

@rflamary @cedricvincentcuaz I haven't stepped through the changes from v0.9.3 to the current state of master (0.9.3...36b4c0a) enough to know if a patch or minor release would be more appropriate, but once PR #630 is in, would it be possible to get a new release with the NumPy 2.0 compatible wheels? I'm not trying to be pushy and I don't want to come across as demanding of your time, so I don't want to ask you to commit to a timeline, I am more just wondering if you think that this would be something that could be done in the near future (~ next week timeline), and if not (understandable as we're all busy) if you have a rough idea of when that might be.

[rflamary]

We discussed with @cedricvincentcuaz and we want a release by the end of the month but next week should be doable. It is OK to push us it means that people need the toolbox ;). I wanted to do a release anyways and this numpy2 stuff is definitely a problem.