ImportError Fixed: Use jax.config.update and remove import .__version

[stevengogogo]

Hello Nico and contributors of tornadox,

Thanks for sharing your work!

When importing your package, I found two issues that cause import error: (1) Importing jax.config is deprecated in newer Jax (2) ._version file is missing. This PR provides a minimal change that required to import with latest Jax version. Please let me know your thoughts!

Issue

The changes of this PR include:

Importing jax.config

Since JAX 0.4.20, the following way to import config is deprecated. This can cause error for newer Jax version, and the current version fails to import

>>> import tornadox
Traceback (most recent call last):
  File "<stdin>", line 1, in <module>
  File "/home/stevenchiu/miniconda3/envs/numjax/lib/python3.11/site-packages/tornadox/__init__.py", line 3, in <module>
    from jax.config import config
ModuleNotFoundError: No module named 'jax.config'

._version not found

Another issue is the ._version is missing in the package. I simply removed this part to make the package imported without error.

Suggestion

As described in this PR, I did two thing in __init__.py

1. Change the way to import jax.config
2. Remove importing ._version

Tested System

jax                       0.4.30                   pypi_0    pypi
jaxlib                    0.4.30                   pypi_0    pypi

Best,
Steven Chiu

[pnkraemer]

Hi, thanks for reaching out!

We don't really work with Tornadox anymore. (We should probably make a note somewhere...).

Instead, I recommend you check out Probdiffeq (https://github.com/pnkraemer/probdiffeq) for JAX code or ProbNumDiffEq.jl (https://github.com/nathanaelbosch/ProbNumDiffEq.jl) for Julia code. These two are more actively maintained. On the JAX side, Probdiffeq can do everything that Tornadox does (and more).

That said, if you would like to work with Tornadox specifically, we can look into this. What do you think? :)

---

Edit: in addition to everything above, the jax.config import change would be helpful! But the version import confuses me slightly. When did you see this error? After installing via pip install or in your fork?

[stevengogogo]

Edit: in addition to everything above, the jax.config import change would be helpful! But the version import confuses me slightly. When did you see this error? After installing via pip install or in your fork?

I used Poetry install, but I agree the verions import might be related to local setup.

I'll look into probdiffeq. For now, I think tornadox provides simple code to learn.

[pnkraemer]

I used Poetry install, but I agree the verions import might be related to local setup.

I agree that it might be a setup issue - the tests pass, which is a great indicator. Thanks for reverting this one change!

I'll look into probdiffeq. For now, I think tornadox provides simple code to learn.

Awesome. I'd love to hear your thoughts once you've played around with these repos a bit :)

[pnkraemer]

Looks good to me.

The linting-related CI fails, but that is unrelated to this PR. I will merge this now.

Thanks for contributing!