T Coxon

I've updated the import to only require `orbax.checkpointing`

Added orbax-checkpoint as a top level import and a package wide dependency. Still have `flax.training` as a non top level import though.

I've just checked, and you are correct, the example doesn't work unless `1` is changed to `1.0`. The specific use case I had envisioned for the current implementation was as...

It has been a long time since I last considered the functionality here. I no longer have an immediate use for this functionality, and as no other users have indicated...

The pre-commit complaints are from bumping the pyright version to `1.1.351`. I've had a quick look in `eqx.Enumeration` to see if this can be handled upstream, but I've drawn a...

I've removed the pyright bump as this wasn't really related to the purpose of this pull.

Thanks for pointing out the potential for silent correctness errors, no doubt this would have caused me some debugging nightmares further down the line! For my case, following your proposed...

That's great and a really clever use of is_leaf (I'd not considered this possibility)! For the scaling problem, if jax (the python xla_client) raised a proper error for the self-referential...

I don't have any other ideas, but personally, I think the hierarchy of types is more straightforward. It avoids the duplicated docstrings, and I think is easier to understand/follow if...