SimonKoop
SimonKoop
The problem disappeared for me after updating JAX at some point. No idea what changed, but I'm happy my code works 😅 Op do 5 sep 2024 om 18:42 schreef...
If you're okay with this, I'd like to work on a pull request for the second option. But if I'm gonna change something about `Sequential`, I do have a question:...
Wow, that was quick! Thanks for fixing this so swiftly!
Why do you need a copy of the parameters? All jittable functions should be without side-effects, and jax Arrays are immutable, so you can just store the original array and...