Hi Patrick,
Firstly, I applaud the excellent paper: Neural SDEs as Infinite-Dimensional GANs.
According to your paper, I am trying to add gradient penalty into your code: /examples/neural_sde.ipynb .
But after my debugging, I found that eqx.filter_grad doesn't seem to accept jax.grad.
I'm a newbie to jax and there's nothing I can do about this. I would be grateful if you could give me some ideas.
Last, Error show:
ValueError: Reverse-mode differentiation does not work for lax.while_loop or lax.fori_loop with dynamic start/stop values. Try using lax.scan, or using fori_loop with static start/stop.
Hi Patrick, Firstly, I applaud the excellent paper: Neural SDEs as Infinite-Dimensional GANs. According to your paper, I am trying to add gradient penalty into your code: /examples/neural_sde.ipynb .
But after my debugging, I found that
eqx.filter_grad
doesn't seem to acceptjax.grad
. I'm a newbie to jax and there's nothing I can do about this. I would be grateful if you could give me some ideas. Last, Error show:ValueError: Reverse-mode differentiation does not work for lax.while_loop or lax.fori_loop with dynamic start/stop values. Try using lax.scan, or using fori_loop with static start/stop.