Allow pass-through implementation of state dict
The previous assert in to_state_dict does not allow to return a state dict that is simply a leaf value. Just returning a leaf is useful for implementing the state dict of box like dataclasses which have just one leaf and some metadata which doesn't end up in the state dict. This way the boxes don't affect the checkpoint format. Before you would get strange looking {"value": x} dictionaries in the checkpoint.
This change just enables the implementation of pass-through state dicts. It doesn't actually change the current implementation of any state dict conversions. So it will not affect existing checkpoints
PiperOrigin-RevId: 560641803
