Open maho3 opened 6 months ago
Hi Matt,
Thanks for pointing out that JAX issue. And sorry for the delay. Have you tried to use the positional argument workaround pointed out there?
I actually ended up implementing the same hack used in your sto
branch and it works for me.
If you want, I can make a PR for this small thing to be put on master.
On the master branch, this line causes the following error:
It seems @eelregit pointed this out as a jax bug, but it hasn't been resolved yet.
Could there be a work around temporarily pushed to master? As this is currently breaking the master branch.