docstr: returns the local shard size and offset on a given tensor dim
docstr_tok: ["returns","the","local","shard","size","and","offset","on","a","given","tensor","dim"]
code: def _local_shard_size_on_dim(
      self,
      size_on_dim: int,
      num_chunks: int,
      rank: int,
      return_offset: bool = False,
  ) -> Tuple[int, int]:
      """
      returns the local shard size and offset on a given tensor dim
      """
      # Compute the chunk size inline with ``torch.chunk``
      full_chunk_size = (size_on_dim + num_chunks - 1) // num_chunks
  
      # Compute chunk size for each chunk on the dimension.
      chunk_sizes = [
          max(
              min(size_on_dim, full_chunk_size * (idx + 1))
              - full_chunk_size * idx,
              0,
          )
          for idx in range(num_chunks)
      ]
      local_shard_size = chunk_sizes[rank]
  
      local_offset_on_dim = -1
      if return_offset:
          # Return global tensor dim size of current dimension if for empty shard
          # to represent the end of the corresponding tensor dim.
          local_offset_on_dim = sum(chunk_sizes[:rank])
  
      return (local_shard_size, local_offset_on_dim)
  
imports: from dataclasses import dataclass
  from typing import Any, cast, List, NamedTuple, Optional, Tuple
  import torch
  import torch.distributed._functional_collectives
  import torch.distributed.distributed_c10d
  from torch.distributed._tensor._collective_utils import mesh_broadcast, mesh_scatter
  from torch.distributed.device_mesh import DeviceMesh
label_desc: Use image node_id 4 for calling the Shard obj's underlying member method code with example usage: obj._local_shard_size_on_dim(size_on_dim, num_chunks, rank, return_offset) and returns: local_shard_size,  local_offset_on_dim