Fixed mypy errors on CI

Author: vfdev-5

flax/traceback_util.py

@@ -25,7 +25,7 @@
 # Whether to filter flax frames from traceback.
 _flax_filter_tracebacks = config.flax_filter_frames
 # Flax specific set of paths to exclude from tracebacks.
-_flax_exclusions = set()
+_flax_exclusions = set()  # type: ignore[var-annotated]
 
 
 # re-import JAX symbol for convenience.

flax/training/checkpoints.py

@@ -1163,7 +1163,7 @@ def read_chunk(i):
     else:
       checkpoint_contents = fp.read()
 
-  state_dict = serialization.msgpack_restore(checkpoint_contents)
+  state_dict = serialization.msgpack_restore(bytes(checkpoint_contents))
   state_dict = _restore_mpas(
     state_dict,
     target,