Flax checkpoint updates

[crstngc]

The focus of this PR is to enable Orbax checkpointing of Flax models and remove an obsolete dependency on Tensorflow.

[codecov]

Codecov Report

Attention: 3 lines in your changes are missing coverage. Please review.

Comparison is base (021d904) 94.58% compared to head (68518e4) 94.80%.

Files	Patch %	Lines
scico/flax/train/state.py	76.92%	3 Missing ⚠️

Additional details and impacted files

@@            Coverage Diff             @@
##             main     #472      +/-   ##
==========================================
+ Coverage   94.58%   94.80%   +0.22%     
==========================================
  Files          90       90              
  Lines        5608     5618      +10     
==========================================
+ Hits         5304     5326      +22     
+ Misses        304      292      -12     

Flag	Coverage Δ
unittests	94.80% <96.34%> (+0.22%)	⬆️

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.

[bwohlberg]

Considering the magnitude of the changes, it would be worth adding an entry in CHANGES.rst.

[bwohlberg]

Will this PR close #409?

[crstngc]

Entries to CHANGES.rst were added. This PR should close #409.

Note that the adaptation of the CLU library implemented to display model parameters remains since the analogous functionality in Flax (based on nn.tabulate) is slightly more complex to deploy and does not include parameter statistics (mean and standard deviations).

[bwohlberg]

Only one minor issue remaining.