[Bug]: Update JAX imports

[cringeyburger]

PyBaMM Version

Develop

Python Version

3.11.5

Describe the bug

While trying solutions for #3617, I discovered that many JAX imports are deprecated. For some reason, this deprecation causes the pytest not to run the tests correctly and pytest only shows deprecation warnings.

I tried fixing a few imports by changing the source code, but it is better to open an issue.

Steps to Reproduce

1. From the file test_base_submodel.py, run any test using pytest (I ran test_parameter_info_error)
2. The code I used to run pytest is as follows:

/home/yukinatsu/miniconda3/bin/conda run -n pybamm --no-capture-output pytest /home/yukinatsu/PyBaMM/tests/unit/test_models/test_submodels/test_base_submodel.py::TestBaseSubModel::test_parameter_info_error

3. You will encounter in total 3 Deprecation Warnings in total (Next warning will come after you solve the previous one)
4. The following are the warnings:

TWO WARNINGS OF : DeprecationWarning: Accessing jax.config via the jax.config submodule is deprecated.

ONE WARNING OF : DeprecationWarning: jax.linear_util.transformation is deprecated. Use jax.extend.linear_util.transformation instead.

Concerned files:

1. jax_bdf_solver.py
2. evaluate_python.py

I changed the imports as needed and there are no further errors and pytest works as intended. I would be happy to take this one.

Relevant log output

No response

[agriyakhetarpal]

Thanks for opening an issue about this @cringeyburger, I'm happy to assign you to it. Looks like I missed some of these deprecation warnings that you have mentioned when I had opened #3644.

[kratman]

@cringeyburger Was this with the latest pybamm? A PR was merged a few hours ago to address the jax_bdf_solver.py warnings (#3671). Just curious if that change was incomplete, if so then feel free to finish it up

[cringeyburger]

Yes, I noticed the recent PR merged by @prady0t, but it seems that a few warnings (solved one was for jax.extend, the jax.config were left out) were left out. I will be opening a PR,