code: def _compute_shards_access(rank_sharding):
      def chunk_offset(sharding):
          assert (
              len(sharding.global_offset)
              == len(sharding.local_shape) + sharding.prepend_axis_num
          )
          return tuple(
              chain(
                  (
                      off
                      for off in sharding.global_offset[
                          : sharding.prepend_axis_num
                      ]
                  ),
                  (
                      off // sh
                      for off, sh in zip(
                          sharding.global_offset[
                              sharding.prepend_axis_num :
                          ],
                          sharding.local_shape,
                      )
                  ),
              )
          )
  
      shard_access_cnt = torch.zeros(
          rank_sharding[0][1].axis_fragmentations,
          dtype=torch.int,
          device="cpu",
      )
      for rank, sharding in rank_sharding:
          if is_main_replica(sharding.replica_id):
              shard_access_cnt[chunk_offset(sharding)] += 1
          # TODO: consider validating different replicas too
      return shard_access_cnt
  
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 12 for calling a global function with example usage: _compute_shards_access(rank_sharding) and returns: shard_access_cnt, tuple