code: def inspect_params(
      *,
      expected_params: PyTree,
      restored_params: PyTree,
      fail_if_extra: bool = True,
      fail_if_missing: bool = True,
      fail_if_shapes_mismatch: bool = False,
  ) -> PyTree:
      """Inspects whether the params are consistent with the expected keys.
  
      Based on
      https://github.com/google-research/big_vision/blob/main/big_vision/model/common.py.
      """
  
      def _flatten_params(d, parent_key="", sep="/"):
          """Flattens a dictionary, keeping empty leaves."""
          items = []
          for k, v in d.items():
              path = parent_key + sep + k if parent_key else k
              if isinstance(v, abc.MutableMapping):
                  items.extend(
                      _flatten_params(v, path, sep=sep).items()
                  )
              else:
                  items.append((path, v))
          # Keeps the empty dict if it was set explicitly.
          if parent_key and not d:
              items.append((parent_key, {}))
          return dict(items)
  
      expected_flat = _flatten_params(
          flax.core.unfreeze(expected_params)
      )
      restored_flat = _flatten_params(
          flax.core.unfreeze(restored_params)
      )
      missing_keys = expected_flat.keys() - restored_flat.keys()
      extra_keys = restored_flat.keys() - expected_flat.keys()
  
      is_shape_mismatch = False
      for key in restored_flat:
          if key in expected_flat:
              restored_shape = None
              expected_shape = None
              # Handle empty nodes (without trainable params)
              if not isinstance(restored_flat[key], dict):
                  restored_shape = restored_flat[key].shape
              if not isinstance(expected_flat[key], dict):
                  expected_shape = expected_flat[key].shape
  
              if restored_shape != expected_shape:
                  is_shape_mismatch = True
                  logging.warning(
                      "Key: %s. Expected shape: %s. Restored shape: %s",
                      key,
                      expected_flat[key].shape,
                      restored_flat[key].shape,
                  )
  
      # Adds back empty dict explicitly, to support layers without weights.
      # Context: FLAX ignores empty dict during serialization.
      empty_keys = set()
      for k in missing_keys:
          if (
              isinstance(expected_flat[k], dict)
              and not expected_flat[k]
          ):
              restored_params[
                  k
              ] = {}  # pytype: disable=unsupported-operands
              empty_keys.add(k)
      missing_keys -= empty_keys
  
      if empty_keys:
          logging.warning(
              "Inspect recovered empty keys:\n%s", empty_keys
          )
  
      logging.info("Inspect missing keys:\n%s", missing_keys)
      logging.info("Inspect extra keys:\n%s", extra_keys)
  
      if fail_if_shapes_mismatch and is_shape_mismatch:
          raise ValueError(
              "Shape mismatch between restored and target model"
          )
  
      if (missing_keys and fail_if_missing) or (
          extra_keys and fail_if_extra
      ):
          raise ValueError(
              f"Missing params from checkpoint: {missing_keys}.\n"
              f"Extra params in checkpoint: {extra_keys}.\n"
              f"Restored params from checkpoint: {restored_flat.keys()}.\n"
              f"Expected params from code: {expected_flat.keys()}."
          )
      return restored_params
  
file_path: scenic/scenic/train_lib/pretrain_utils.py
imports: from collections import abc
  import os
  import re
  from typing import Any, Dict, Mapping, List, Optional, Union
  from absl import logging
  import flax
  from flax.training import checkpoints
  import numpy
  from scenic.train_lib import train_utils
  from tensorflow.io import gfile
label_desc: Use image node_id 4 for calling a global function with example usage: inspect_params() and returns: restored_params, dict