Hi,
You question is a bit confusing between which implementations you have and which match and which don’t.
The example you gave works fine for me.
What most likely happens is that you hit a point of non-differentiability of your function and that torch and jax behave differently at this point.
Can you check what is the point where your function becomes nan? Also you can try to enable the anomaly mode in pytorch to see which op created the nans.