Migrate Flax dataclass to the newest JAX pytree keypath API.
* `flax.struct.dataclass` already registered via `register_keypath`, and this one only changes it to latest API.
* `flax.core.FrozenDict` was registered so that flattening a frozen dict should be the same as flattening the underlying dict. This makes its serialization backward-compatible.
PiperOrigin-RevId: 515096454
