Parameter update of a Neural Network using JAX

[raj-brown]

Dear All-
Is there any example for updating the neural network parameters using jax.jvp with simple sgd approach?

Thanks!