Flax dataclasses do not support new-style PRNGKeyArray in serialisation

New-style `PRNGKeyArray` cannot be serialised by flax. 
This is inconveniente as one has to manually override the serailization/deserialization methods automatically generated by flax's data class to wrap/unwrap those fields..

Would it be possible to support natively PRNGKeyArray (which are technically leafs) within flax.serialization? 

