Alex

Ok. I have the same issue. ``` python train.py --problem=mnist --save_path=/tmp/mnist_opt Traceback (most recent call last): File "train.py", line 28, in import meta File "/Volumes/ExtraDiskSpace/Documents/learning-to-learn/meta.py", line 32, in import networks...

I also want a div op. It seems the authors of `local_kinetic_energy` in [here](https://github.com/deepmind/ferminet/blob/jax/ferminet/hamiltonian.py) implement a div operator, kinda (to compute the laplacian). Also. On the topic of shortcuts for...

I tried with ot.emd2. Same issue. Like I said, the issue isn't with grad. It also applies to jit, vmap, ... any of jax's code transformation fns. In the ot...