ImportError: cannot import name 'isin' from 'jax._src.numpy.lax_numpy'

[FelipeMartins96]

Following the getting started steps for installing dopamine from source on a python3.10 virtual environment, I got the following error:

ImportError: cannot import name 'isin' from 'jax._src.numpy.lax_numpy'

this stackoverflow question has information on the same problem; in my environment, the flax version installed was 0.3.6, and updating it to the most recent (0.4.2) solved this issue. However, I could not attempt further tests if only updating the flax version o requirements is possible