code: def validate_shards(
      shards: Sequence[TensorSliceDict],
      shardable_tensors: Sequence[ShardableTensor],
      callback_description: str,
  ) -> None:
      """Validates shards generated by the sharding_callback."""
      unseen_tensor_dict = {}
      for shardable_tensor in shardable_tensors:
          unseen_tensor_dict.setdefault(
              shardable_tensor.checkpoint_key, {}
          )[shardable_tensor.slice_spec] = shardable_tensor.tensor
      seen_tensor_set = set()
  
      for shard_tensors in shards:
          task_tensor = None
          for (
              checkpoint_key,
              tensor_slice_dict,
          ) in shard_tensors.items():
              for slice_spec, shard_tensor in tensor_slice_dict.items():
                  slice_spec = slice_spec.strip()
  
                  # Validate uniqueness.
                  if (checkpoint_key, slice_spec) in seen_tensor_set:
                      raise RuntimeError(
                          "After executing the checkpoint sharding callback, multiple "
                          "tensors with the same checkpoint key and slice spec were "
                          "found:\n"
                          f"  callback_description: {callback_description}\n"
                          f"  checkpoint_key: {checkpoint_key}\n"
                          f"  slice_spec: {slice_spec}\n"
                      )
  
                  # Validate no added tensors.
                  if checkpoint_key not in unseen_tensor_dict:
                      raise RuntimeError(
                          "After executing the checkpoint sharding callback, a tensor "
                          "not originally in the object graph was found in the "
                          "checkpoint shards:\n"
                          f"  callback_description: {callback_description}\n"
                          f"  checkpoint_key: {checkpoint_key}\n"
                          f"  slice_spec: {slice_spec}\n"
                      )
  
                  # Validate no shape change.
                  target_shape = unseen_tensor_dict[checkpoint_key][
                      slice_spec
                  ].shape
                  if shard_tensor.shape != target_shape:
                      raise RuntimeError(
                          "After executing the checkpoint sharding callback, a tensor "
                          "was found with an altered shape:\n"
                          f"  callback_description: {callback_description}\n"
                          f"  checkpoint_key: {checkpoint_key}\n"
                          f"  slice_spec: {slice_spec}\n"
                          f"  original tensor_shape: {target_shape}\n"
                          f"  new tensor_shape: {shard_tensor.shape}\n"
                      )
  
                  # Validate no dtype change.
                  target_dtype = unseen_tensor_dict[checkpoint_key][
                      slice_spec
                  ].dtype
                  if shard_tensor.dtype != target_dtype:
                      raise RuntimeError(
                          "After executing the checkpoint sharding callback, a tensor "
                          "was found with an altered dtype:\n"
                          f"  callback_description: {callback_description}\n"
                          f"  checkpoint_key: {checkpoint_key}\n"
                          f"  slice_spec: {slice_spec}\n"
                          f"  original tensor_dtype: {target_dtype}\n"
                          f"  new tensor_dtype: {shard_tensor.dtype}\n"
                      )
  
                  # Validate same task in shard.
                  if task_tensor is None:
                      task_tensor = ShardableTensor
                      task_tensor.device = shard_tensor.device
                      task_tensor.checkpoint_key = checkpoint_key
                      task_tensor.slice_spec = slice_spec
                  else:
                      task1 = device_lib.DeviceSpec.from_string(
                          task_tensor.device
                      ).task
                      task2 = device_lib.DeviceSpec.from_string(
                          shard_tensor.device
                      ).task
                      if (
                          task1 is not None
                          and task2 is not None
                          and task1 != task2
                      ):
                          raise RuntimeError(
                              "After executing the checkpoint sharding callback, tensors "
                              "with different tasks were found in the same shard:\n"
                              f"  callback_description: {callback_description}\n"
                              "  tensor #1:"
                              f"    checkpoint_key: {task_tensor.checkpoint_key}\n"
                              f"    slice_spec: {task_tensor.slice_spec}\n"
                              f"    task: {task1}\n"
                              "  tensor #2:"
                              f"    checkpoint_key: {checkpoint_key}\n"
                              f"    slice_spec: {slice_spec}\n"
                              f"    task: {task2}\n"
                          )
  
                  del unseen_tensor_dict[checkpoint_key][slice_spec]
                  if not unseen_tensor_dict[checkpoint_key]:
                      del unseen_tensor_dict[checkpoint_key]
                  seen_tensor_set.add((checkpoint_key, slice_spec))
  
      # validate no tensor removal
      if unseen_tensor_dict:
          tensors_info = ""
          for ckpt_key, slice_spec in unseen_tensor_dict.items():
              tensors_info += "  tensor:\n"
              tensors_info += f"    checkpoint_key: {ckpt_key}\n"
              tensors_info += f"    slice_spec: {slice_spec}\n"
          raise RuntimeError(
              "After executing the checkpoint sharding callback, tensors in the "
              "object graph were not found in the checkpoint shards:\n"
              f"  callback_description: {callback_description}\n"
              f"{tensors_info}"
          )
  
imports: import abc
  import dataclasses
  import inspect
  from typing import Hashable, MutableMapping, Sequence
  from tensorflow.python.framework import device
  from tensorflow.python.framework import dtypes
  from tensorflow.python.framework import tensor
  from tensorflow.python.framework import tensor_shape
  from tensorflow.python.framework import tensor_spec
  from tensorflow.python.ops import variables
  from tensorflow.python.trackable import base
  from tensorflow.python.training.saving import saveable_object
  from tensorflow.python.util import tf_export
label_desc: Use image node_id 1 for calling a global function with example usage: validate_shards(shards, shardable_tensors, callback_description) without return types