I think this is my bad. I used the regular optimizer.SR in the example instead of that for the Jax machines optimizer.JaxSR. Running it with the optimizer.JaxSR works for me.
Ah thanks Tom! Indeed I didn't realize that... I fixed it now (even though the dispatch over the "right" SR will be automatic in the future, so the example is actually OK as it is)
