Error on trying to run the experiments about easy-neural-ode

[jacobjinkelly]

Hi @avik-pal, thanks for making this issue! I'm a little busy right now but I can followup this upcoming weekend

[avik-pal]

@jacobjinkelly yes it works now. Thanks a lot!

[avik-pal]

cc @jacobjinkelly @jessebett

[avik-pal]

@jacobjinkelly were you able to figure this out?

[jacobjinkelly]

Hi @avik-pal, sorry for dropping the ball on this! Could you try jax==0.1.64 and jaxlib==0.1.45? I'm basing this off of the original code we used for the experiments which is in this fork of jax. If that doesn't work, let me know and I'm happy to dig deeper.

[qsh-zh]

@avik-pal I encounter the same errors. Do you figure this out? @jacobjinkelly Thanks for your open source project. However, changing the version of jax does not work. Could you help report your environment setup?

[jacobjinkelly]

Hi @avik-pal @qshzh I managed to figure it out, all of the fixes are pushed. I got the code to run with the latest version of JAX. Let me know if that works for you!