nice error

Author: khatwanimohit

src/maxtext/utils/model_creation_utils.py

@@ -345,6 +345,15 @@ def create_sharded_state():
         else:
           checkpoint = restored["params"]["params"]
 
+        loaded_count = len(jax.tree_util.tree_leaves(checkpoint))
+        expected_count = len(jax.tree_util.tree_leaves(target_for_restore))
+        if loaded_count < expected_count:
+          raise ValueError(
+              f"Checkpoint at '{config.load_parameters_path}' loaded only {loaded_count} of {expected_count} "
+              "expected parameter arrays. This usually means a scanned (stacked-layers) checkpoint was provided "
+              "where an unscanned checkpoint is required. Please convert the checkpoint to unscanned format first."
+          )
+
         if checkpoint:
           model_arrays = jax.tree.map(
               lambda v: v.value,