WIP: Autograd jax example conversion

At @tisaac's request, this MR exists to track our progress on converting the existing autograd examples to jax.

• The p-Stokes example has been converted to JAX in !16 (merged).
• The p-Laplacian example still uses Autograd and TensorFlow.

(This is a re-opening of !22 (merged)).