Support Python 3.13

[elliotwutingfeng]

Description

Update all relevant references to Python version from 3.12 to 3.13, in documentation and CI.

Fixes #4539

Type of change

• New feature (non-breaking change which adds functionality)
• Optimization (back-end change that speeds up the code)
• Bug fix (non-breaking change which fixes an issue)

Key checklist:

• No style issues: $ pre-commit run (or $ nox -s pre-commit) (see CONTRIBUTING.md for how to set this up to run automatically when committing locally, in just two lines of code)
• All tests pass: $ python run-tests.py --all (or $ nox -s tests)
• The documentation builds: $ python run-tests.py --doctest (or $ nox -s doctests)

You can run integration tests, unit tests, and doctests together at once, using $ python run-tests.py --quick (or $ nox -s quick).

Further checks:

• Code is commented, particularly in hard-to-understand areas
• Tests added that prove fix is effective or that feature works

[elliotwutingfeng]

Looks like the old jaxlib version is incompatible with Python 3.13.

[kratman]

Looks like the old jaxlib version is incompatible with Python 3.13.

I think newer versions of Jax also drop support for python 3.9, so we might have to drop support of 3.9 in favor of 3.13

[kratman]

What you need to do next is try to find a Jax version that supports both version 3.9 and version 3.13 for python. It needs to be updated in both utils and the pyproject file.

If you can't find a version that supports both, then we have to do a few other tasks first:

• Remove 3.9
• Remove IREE support (this is planned and is likely needed even if you find a working version of Jax)

I can help you out if needed since this somewhat overlaps with my current tasks

[elliotwutingfeng]

I think I'd need help on this one!

[agriyakhetarpal]

I think I'd need help on this one!

Sure, what can we help you with, @elliotwutingfeng? Based on the suggestions provided above, JAX wheels for Python 3.13 are available post version 0.4.34; you can support multiple versions of JAX using specifiers, i.e., with something like this in the relevant section of the dependencies table:

# For the Jax solver.
# Note: These must be kept in sync with the versions defined in pybamm/util.py, and
#       must remain compatible with IREE (see noxfile.py for IREE compatibility).
jax = [
    "jax==0.4.27; python_version <'3.13'",
    "jax==0.4.34; python_version <'3.13'"
]

With this, we don't necessarily need to drop Python 3.9 support right now. However, there's no guarantee that version 0.4.34 will pass the non-IREE (or IREE, even) tests, so you'll need to fix any failing tests if necessary so that both version 0.4.27 and 0.4.34 work (it's unlikely, given how much of JAX's internals we use and the breaking changes there are present between their releases – supporting both is more of a code smell). We can wait on this until we decide on IREE and on updating the JAX version, or if you feel ambitious, please go ahead. cc @jsbrittain who has been working on retaining the IREE compiler support with advancing JAX versions just in case he has anything to comment on.

[codecov]

Codecov Report

All modified and coverable lines are covered by tests ✅

Project coverage is 99.26%. Comparing base (bcdd0b5) to head (7aa4ef0).

Additional details and impacted files

@@           Coverage Diff            @@
##           develop    #4552   +/-   ##
========================================
  Coverage    99.26%   99.26%           
========================================
  Files          302      302           
  Lines        22889    22890    +1     
========================================
+ Hits         22721    22722    +1     
  Misses         168      168           

☔ View full report in Codecov by Sentry.
📢 Have feedback on the report? Share it here.

[kratman]

@elliotwutingfeng Some of your failures can be solved by updating the IREE state function in the nox file:

def set_iree_state():
    """
    Check if IREE is enabled and set the environment variable accordingly.

    Returns
    -------
    str
        "ON" if IREE is enabled, "OFF" otherwise.

    """
    state = "ON" if os.getenv("PYBAMM_IDAKLU_EXPR_IREE", "OFF") == "ON" else "OFF"
    if state == "ON":
        if sys.platform == "win32" or sys.platform == "darwin":
            warnings.warn(
                (
                    "IREE is not enabled on Windows and MacOS. "
                    "Setting PYBAMM_IDAKLU_EXPR_IREE=OFF."
                ),
                stacklevel=2,
            )
            state = "OFF"
        if sys.version_info >= (3, 13):
            warnings.warn(
                (
                    "IREE is not available for Python 3.13 or higher. "
                    "Setting PYBAMM_IDAKLU_EXPR_IREE=OFF."
                ),
                stacklevel=2,
            )
            state = "OFF"
    return state

This is part of what I was saying before about this impacting both IREE and 3.9 support. Both of those should be removed in the next year or so.

[kratman]

@elliotwutingfeng Are you still working on this or do you need some help still?

[elliotwutingfeng]

Apologies for the delayed response, I've been caught up with other commitments over the past few days. I think it'll be best for your team to take over this PR. 😅

[kratman]

Apologies for the delayed response, I've been caught up with other commitments over the past few days. I think it'll be best for your team to take over this PR. 😅

I was not trying to rush you, just asking since we have a release coming out soon. I think you are only 1-2 commits away from being done, but I am more than happy to help you wrap this up

[kratman]

I will try work on that windows failure this weekend. It could be that we just need to update numpy before that works

[kratman]

I tried this out locally on a windows machine with python 3.13, and it looked like it was trying to build some libraries from source for numpy. This was a different failure than the one in the CI (a warning about MinGW built numpy being unstable and likely to crash, right before crashing)

I think this will require updating numpy to get this all working. I am taking a look at that to see how much work would be required