Flax will default to returning regular dicts instead of FrozenDicts i…

…n the future. Some internal-only tests fail because they expect frozen params when creating TrainState.

PiperOrigin-RevId: 533567519

flax/training/train_state.py

@@ -83,11 +83,12 @@ def apply_gradients(self, *, grads, **kwargs):
   @classmethod
   def create(cls, *, apply_fn, params, tx, **kwargs):
     """Creates a new instance with `step=0` and initialized `opt_state`."""
-    opt_state = tx.init(params)
+    frozen_params = core.freeze(params) if isinstance(params, dict) else params
+    opt_state = tx.init(frozen_params)
     return cls(
         step=0,
         apply_fn=apply_fn,
-        params=params,
+        params=frozen_params,
         tx=tx,
         opt_state=opt_state,
         **kwargs,