[FEATURE] Jax tutorial notebook

[mmorris44]

Please describe the purpose of the feature. Is it related to a problem?

A clear and concise description of what the problem is. Ex. I'm always frustrated when [...]

Describe the solution you'd like

A clear and concise description of what you want to happen.

Describe alternatives you've considered

A clear and concise description of any alternative solutions or features you've considered.

How do we know when implementation of this feature is complete?

Checklist:

• [ ]
• [ ]

Additional context

Add any other context about the feature request here.

TODO: look at PyTorch lightning for inspiration