What version of jax is used here?

[martavarela]

It is hard to get the code running without that information (or the version of Python and other packages used).

Thank you, Marta