[FEATURE] Jax Checkpointers

[mmorris44]

Feature

Custom checkpointers for jax systems

Proposal

@RuanJohn found a nice easy way to do this.

Testing

Unique tests will need to be created. Checkpointers should also be tested on an example.

Definition of done

Checkpointers component is created, working in example, and tested.

Mandatory checklist before making a PR

• The success criteria laid down in “Definition of done” are met.
• Code is documented - docstrings for methods and classes, static types for arguments.
• Code is tested - unit, integration and/or functional tests are added.
• Documentation is updated - README, CONTRIBUTING, or other documentation.
• All functional tests are green.
• Link experiment/benchmarking after implementation (optional).