CustomVJPException plus memory leak when using a for loop instead of a scan.

I've been using `jaxopt.implicit_diff.custom_root` for differentiating through an jax-md energy minimization routine and I have noticed that if I am using a python for loop for my solver then I get a `CustomVJPException` and an additional memory leak.

This memory leak only seems to show up when I get the `CustomVJPException` and not when I modify my code to prevent the exception from happening. I believe the underlying reason for that exception is the same as in issue #31 and seems to stem from the fact how jax-md defines its energy functions.

I'd like to know how to change that part of jax-md to prevent the `CustomVJPException` from happening in the first place but I've haven't managed to come up with a simplified version that would let me pinpoint the source of the error. But I can give it another shot if that helps you.

Here's a colab demo that demonstrates the issue.
https://colab.research.google.com/drive/1f_3EmFQpvW1p7A1AcNw8uqX5T79fjXRS?usp=sharing