Add checkpointing callbacks

[mj-will]

This PR adds support for a checkpoint callback: a function that is called instead of the default checkpoint function when the sampler tries to checkpoint.

This is motivated by the discussion in gwastro/pycbc#4567

Usage

def callback(state):
    # Save the state
    ...

fs = FlowSampler(
    model,
    ...,
    checkpoint_callback=callback,
)
fs.run()

To-Do

• Documentation

[codecov]

Codecov Report

All modified and coverable lines are covered by tests ✅

Comparison is base (5d93ca2) 93.26% compared to head (c96ecb9) 93.26%.

Additional details and impacted files

@@           Coverage Diff           @@
##             main     #355   +/-   ##
=======================================
  Coverage   93.26%   93.26%           
=======================================
  Files          67       67           
  Lines        6296     6300    +4     
=======================================
+ Hits         5872     5876    +4     
  Misses        424      424           

Flag	Coverage Δ
unittests	93.26% <100.00%> (+<0.01%)	⬆️

Flags with carried forward coverage won't be shown. Click here to find out more.

☔ View full report in Codecov by Sentry.
📢 Have feedback on the report? Share it here.