code: def test_shard_checkpoint(self):
      # This is the model we will use, total size 340,000 bytes.
      model = tf.keras.Sequential(
          [
              tf.keras.layers.Dense(200, use_bias=False),  # size 80,000
              tf.keras.layers.Dense(
                  200, use_bias=False
              ),  # size 160,000
              tf.keras.layers.Dense(100, use_bias=False),  # size 80,000
              tf.keras.layers.Dense(50, use_bias=False),  # size 20,000
          ]
      )
      inputs = tf.zeros((1, 100), dtype=tf.float32)
      model(inputs)
      weights = model.weights
      weights_dict = {w.name: w for w in weights}
      with self.subTest(
          "No shard when max size is bigger than model size"
      ):
          shards, index = tf_shard_checkpoint(weights)
          self.assertIsNone(index)
          self.assertDictEqual(shards, {TF2_WEIGHTS_NAME: weights})
  
      with self.subTest(
          "Test sharding, no weights bigger than max size"
      ):
          shards, index = tf_shard_checkpoint(
              weights, max_shard_size="300kB"
          )
          # Split is first two layers then last two.
          self.assertDictEqual(
              index,
              {
                  "metadata": {"total_size": 340000},
                  "weight_map": {
                      "dense/kernel:0": "tf_model-00001-of-00002.h5",
                      "dense_1/kernel:0": "tf_model-00001-of-00002.h5",
                      "dense_2/kernel:0": "tf_model-00002-of-00002.h5",
                      "dense_3/kernel:0": "tf_model-00002-of-00002.h5",
                  },
              },
          )
  
          shard1 = [
              weights_dict["dense/kernel:0"],
              weights_dict["dense_1/kernel:0"],
          ]
          shard2 = [
              weights_dict["dense_2/kernel:0"],
              weights_dict["dense_3/kernel:0"],
          ]
          self.assertDictEqual(
              shards,
              {
                  "tf_model-00001-of-00002.h5": shard1,
                  "tf_model-00002-of-00002.h5": shard2,
              },
          )
  
      with self.subTest(
          "Test sharding with weights bigger than max size"
      ):
          shards, index = tf_shard_checkpoint(
              weights, max_shard_size="100kB"
          )
          # Split is first layer, second layer then last 2.
          self.assertDictEqual(
              index,
              {
                  "metadata": {"total_size": 340000},
                  "weight_map": {
                      "dense/kernel:0": "tf_model-00001-of-00003.h5",
                      "dense_1/kernel:0": "tf_model-00002-of-00003.h5",
                      "dense_2/kernel:0": "tf_model-00003-of-00003.h5",
                      "dense_3/kernel:0": "tf_model-00003-of-00003.h5",
                  },
              },
          )
  
          shard1 = [weights_dict["dense/kernel:0"]]
          shard2 = [weights_dict["dense_1/kernel:0"]]
          shard3 = [
              weights_dict["dense_2/kernel:0"],
              weights_dict["dense_3/kernel:0"],
          ]
          self.assertDictEqual(
              shards,
              {
                  "tf_model-00001-of-00003.h5": shard1,
                  "tf_model-00002-of-00003.h5": shard2,
                  "tf_model-00003-of-00003.h5": shard3,
              },
          )
  
imports: from __future__ import annotations
  import inspect
  import json
  import os
  import random
  import tempfile
  import unittest
  import unittest.mock
  from huggingface_hub import HfFolder, Repository, delete_repo, snapshot_download
  from huggingface_hub.file_download import http_get
  from requests.exceptions import HTTPError
  from transformers import is_tf_available, is_torch_available
  from transformers.configuration_utils import PretrainedConfig
  from transformers.testing_utils import TOKEN, USER, CaptureLogger, _tf_gpu_memory_limit, is_pt_tf_cross_test, is_staging_test, require_safetensors, require_tf, require_torch, slow
  from transformers.utils import SAFE_WEIGHTS_NAME, TF2_WEIGHTS_INDEX_NAME, TF2_WEIGHTS_NAME, logging
label_desc: Use image node_id 12 for calling the TFModelUtilsTest obj's underlying member method code with example usage: obj.test_shard_checkpoint() without return types