code: def test_checkpoint_sharding_local(self):
      model = FlaxBertModel.from_pretrained(
          "hf-internal-testing/tiny-bert-flax-only"
      )
  
      with tempfile.TemporaryDirectory() as tmp_dir:
          # We use the same folder for various sizes to make sure a new save erases the old checkpoint.
          for max_size in ["150kB", "150kiB", "200kB", "200kiB"]:
              model.save_pretrained(tmp_dir, max_shard_size=max_size)
  
              # Get each shard file and its size
              shard_to_size = {}
              for shard in os.listdir(tmp_dir):
                  if shard.endswith(".msgpack"):
                      shard_file = os.path.join(tmp_dir, shard)
                      shard_to_size[shard_file] = os.path.getsize(
                          shard_file
                      )
  
              index_file = os.path.join(
                  tmp_dir, FLAX_WEIGHTS_INDEX_NAME
              )
              # Check there is an index but no regular weight file
              self.assertTrue(os.path.isfile(index_file))
              self.assertFalse(
                  os.path.isfile(
                      os.path.join(tmp_dir, FLAX_WEIGHTS_NAME)
                  )
              )
  
              # Check a file is bigger than max_size only when it has a single weight
              for shard_file, size in shard_to_size.items():
                  if max_size.endswith("kiB"):
                      max_size_int = int(max_size[:-3]) * 2**10
                  else:
                      max_size_int = int(max_size[:-2]) * 10**3
                  # Note: pickle adds some junk so the weight of the file can end up being slightly bigger than
                  # the size asked for (since we count parameters)
                  if size >= max_size_int + 50000:
                      with open(shard_file, "rb") as state_f:
                          state_file = from_bytes(
                              FlaxBertModel, state_f.read()
                          )
                          self.assertEqual(len(state_file), 1)
  
              # Check the index and the shard files found match
              with open(index_file, "r", encoding="utf-8") as f:
                  index = json.loads(f.read())
  
              all_shards = set(index["weight_map"].values())
              shards_found = set(
                  f
                  for f in os.listdir(tmp_dir)
                  if f.endswith(".msgpack")
              )
              self.assertSetEqual(all_shards, shards_found)
  
              # Finally, check the model can be reloaded
              new_model = FlaxBertModel.from_pretrained(tmp_dir)
              for p1, p2 in zip(
                  flatten_dict(model.params).values(),
                  flatten_dict(new_model.params).values(),
              ):
                  self.assertTrue(
                      np.allclose(np.array(p1), np.array(p2))
                  )
  
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 30 for calling the FlaxModelTesterMixin obj's underlying member method code with example usage: obj.test_checkpoint_sharding_local() without return types