code: def get_shard_endpoints(
      dim_size: int,
      shard_rank: int,
      num_shards: int,
      granularity: int = DEFAULT_SHARD_GRANULARITY,
  ) -> Tuple[int, int]:
      """
      Given a dimension to shard with size dim_size, return the start and end indices of the slice
      that belong to the given rank.
  
      The typical use of this is as an internal helper function, so see if there is a higher level
      API that better suits the application.
  
      Args:
          dim_size (int): The size of the dimension to shard.
          shard_rank (int): The rank of the shard to return.
          num_shards (int): Total number of shards the dimension will be distributed across.
          granularity (int): The minimum alignment of the shard endpoints. This is used to support
              non-even head counts as well as align dimensions to cleaner GEMM boundaries.
      """
      assert (
          dim_size % granularity == 0
      ), "Dimension size must be divisible by granularity"
  
      total_chunks = dim_size // granularity
      base_chunks_per_rank = total_chunks // num_shards
      remainder_chunks = total_chunks % num_shards
  
      start_chunk_id = shard_rank * base_chunks_per_rank + min(
          shard_rank, remainder_chunks
      )
      end_chunk_id = (
          start_chunk_id
          + base_chunks_per_rank
          + (1 if shard_rank < remainder_chunks else 0)
      )
  
      return start_chunk_id * granularity, end_chunk_id * granularity
  
imports: from typing import Optional, Tuple
  import torch
  from .types import ShardingType, DEFAULT_SHARD_GRANULARITY
label_desc: Use image node_id 1 for calling a global function with example usage: get_shard_endpoints(dim_size, shard_rank, num_shards, granularity) and returns: unknown,  unknown