code: def test_save_load_from_base(self):
      (
          config,
          _,
      ) = self.model_tester.prepare_config_and_inputs_for_common()
      base_class = FLAX_MODEL_MAPPING[config.__class__]
  
      for model_class in self.all_model_classes:
          if model_class == base_class:
              continue
  
          model = base_class(config)
          base_params = flatten_dict(unfreeze(model.params))
  
          # check that all base model weights are loaded correctly
          with tempfile.TemporaryDirectory() as tmpdirname:
              model.save_pretrained(tmpdirname)
              head_model = model_class.from_pretrained(tmpdirname)
  
              base_param_from_head = flatten_dict(
                  unfreeze(
                      head_model.params[head_model.base_model_prefix]
                  )
              )
  
              for key in base_param_from_head.keys():
                  max_diff = (
                      (base_params[key] - base_param_from_head[key])
                      .sum()
                      .item()
                  )
                  self.assertLessEqual(
                      max_diff, 1e-3, msg=f"{key} not identical"
                  )
  
imports: import copy
  import inspect
  import json
  import random
  import tempfile
  import unittest
  from typing import List, Tuple
  import numpy
  import transformers
  from huggingface_hub import HfFolder, delete_repo, set_access_token
  from requests.exceptions import HTTPError
  from transformers import BertConfig, is_flax_available, is_torch_available
  from transformers.models.auto import get_values
  from transformers.testing_utils import TOKEN, USER, CaptureLogger, is_pt_flax_cross_test, is_staging_test, require_flax, torch_device
  from transformers.utils import logging
  from transformers.utils.generic import ModelOutput
label_desc: Use image node_id 8 for calling the FlaxModelTesterMixin obj's underlying member method code with example usage: obj.test_save_load_from_base() without return types