GitInTheRobot

> What does the `doJax` function contain? It calls some flax modules to train a neural network. I posted here because I assumed that memory deallocation is handled/handle-able through JAX.

I'll see if I can create it in the future, but it is a fairly large, branching set of code so it will take some time to distill the problematic...