Compatibility with `jax==0.6.0`

[f-dangel]

Hi,

I'm trying to use folx (0.2.15) with the latest JAX (0.6.0), but get an import error when I try

import folx

that looks as follows:

[...]
    from jax.core import Primitive
ImportError: cannot import name 'Primitive' from 'jax.core'

Could you tell me what version I need to downgrade JAX to, or fix this import error in a new release?

Thanks a lot,
Felix

Quick update: Downgrading to jax==0.5.3 fixed the import error.

[n-gao]

Hi, folx surely works <0.5, I'll have a look at the changes in Jax at a later point.

[n-gao]

@f-dangel I just released folx>=0.2.26 which supports all jax>=0.4.10.

[f-dangel]

Lovely, thanks a lot for the quick fix!