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.27%. Comparing base (f05cae2) to head (50618c3).

Additional details and impacted files

@@           Coverage Diff            @@
##           develop    #4552   +/-   ##
========================================
  Coverage    99.27%   99.27%           
========================================
  Files          300      300           
  Lines        22795    22796    +1     
========================================
+ Hits         22630    22631    +1     
  Misses         165      165           

☔ 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.