code: def _validate_sharding_for_key_flattened(tensors_by_shard):
      all_slices = []
      local_shape = tensors_by_shard[0].local_shape
      for sharding in tensors_by_shard:
          assert sharding.local_shape == local_shape
          sharding: ShardedTensor
          if not is_main_replica(sharding.replica_id):
              # TODO: this checks only saving (and loading replica_id=0) consistency
              continue
  
          all_slices.append(
              (
                  sharding.flattened_range.start,
                  sharding.flattened_range.stop,
              )
          )
  
      starts, stops = map(np.asarray, zip(*sorted(all_slices)))
      if (
          starts[0] != 0
          or stops[-1] != np.product(local_shape)
          or not np.all(starts[1:] == stops[:-1])
      ):
          logger.error(
              f"Flattened ranges dont cover the whole shard {tensors_by_shard[0]}. Ranges: {(starts, stops)}"
          )
          raise CheckpointingException(
              f"Flattened ranges dont cover the whole shard {tensors_by_shard[0]}"
          )
  
imports: import logging
  import os
  from collections import Counter, defaultdict
  from itertools import chain
  from pathlib import Path
  from typing import Iterable, List, Optional, Tuple, Union
  import numpy
  import torch
  from .core import CheckpointingConfig, maybe_load_config, save_config
  from .dict_utils import dict_list_map_inplace, diff, extract_matching_values, map_reduce, merge, nested_values
  from .mapping import CheckpointingException, ShardedObject, ShardedStateDict, ShardedTensor, ShardedTensorFactory, StateDict, apply_factories, apply_factory_merges, is_main_replica
  from .strategies.base import LoadCommonStrategy, LoadShardedStrategy, SaveCommonStrategy, SaveShardedStrategy, StrategyAction, get_default_strategy
  from .utils import extract_sharded_tensors, extract_sharded_tensors_or_nonpersistent
label_desc: Use image node_id 13 for calling a global function with example usage: _validate_sharding_for_key_flattened(tensors_by_shard) without return types