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[bot]]

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 :warning:

Additional details and impacted files

```diff @@ 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](https://app.codecov.io/gh/lanl/scico/pull/472/flags?src=pr&el=flags&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=lanl) | Coverage Δ | | |---|---|---| | [unittests](https://app.codecov.io/gh/lanl/scico/pull/472/flags?src=pr&el=flag&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=lanl) | `94.80% <96.34%> (+0.22%)` | :arrow_up: | Flags with carried forward coverage won't be shown. [Click here](https://docs.codecov.io/docs/carryforward-flags?utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=lanl#carryforward-flags-in-the-pull-request-comment) to find out more.

:umbrella: View full report in Codecov by Sentry.
:loudspeaker: Have feedback on the report? Share it here.

[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).