docstr: This function uses torch.chunk to split a tensor into num_chunks shards along
  the Shard placement dimension, and return a list of shards with their pad sizes.
  
  Keyword args:
      with_padding (bool, optional): when True, we pad the tensor on the last
      few ranks before calling the collectives (i.e. scatter/all_gather, etc.).
      This is because collectives usually require equal size tensor inputs
docstr_tok: ["This","function","uses","torch.chunk","to","split","a","tensor","into","num_chunks","shards","along","the","Shard","placement","dimension",",","and","return","a","list","of","shards","with","their","pad","sizes",".","Keyword","args",":","with_padding","(","bool",",","optional",")",":","when","True",",","we","pad","the","tensor","on","the","last","few","ranks","before","calling","the","collectives","(","i.e",".","scatter\/all_gather",",","etc.",")",".","This","is","because","collectives","usually","require","equal","size","tensor","inputs"]
code: def _split_tensor(
      self,
      tensor: torch.Tensor,
      num_chunks: int,
      *,
      with_padding: bool = True,
      contiguous: bool = True,
  ) -> Tuple[List[torch.Tensor], List[int]]:
      """
      This function uses torch.chunk to split a tensor into num_chunks shards along
      the Shard placement dimension, and return a list of shards with their pad sizes.
  
      Keyword args:
          with_padding (bool, optional): when True, we pad the tensor on the last
          few ranks before calling the collectives (i.e. scatter/all_gather, etc.).
          This is because collectives usually require equal size tensor inputs
      """
      assert (
          self.dim <= tensor.ndim
      ), f"Sharding dim {self.dim} greater than tensor ndim {tensor.ndim}"
  
      # chunk tensor over dimension `dim` into n slices with padding if necessary
      tensor_list = list(torch.chunk(tensor, num_chunks, dim=self.dim))
      # compute the chunk size inline with ``torch.chunk``
      full_chunk_size = (
          tensor.size(self.dim) + num_chunks - 1
      ) // num_chunks
  
      # Compute chunk size for each chunk for ``self.dim``
      chunk_sizes = [
          tensor_list[idx].size(self.dim)
          if idx < len(tensor_list)
          else 0
          for idx in range(num_chunks)
      ]
      # Compute pad size on each chunk
      pad_sizes = [
          full_chunk_size - chunk_size for chunk_size in chunk_sizes
      ]
  
      # Reuse tensor to fill empty chunk with empty tensor
      num_empty_tensors = num_chunks - len(tensor_list)
      tensor_size = list(tensor_list[0].size())
      tensor_size = [
          size if idx != self.dim else 0
          for idx, size in enumerate(tensor_size)
      ]
      tensor = tensor.new_zeros(tensor_size)
      for _ in range(num_empty_tensors):
          tensor_list.append(tensor)
  
      if with_padding or contiguous:
          shard_list = []
          for shard, pad_size in zip(tensor_list, pad_sizes):
              # Fill the empty tensor with zeroes with padding.
              if with_padding and pad_size > 0:
                  shard = self._pad_tensor(shard, pad_size)
              shard = shard.contiguous() if contiguous else shard
              shard_list.append(shard)
          return shard_list, pad_sizes
      else:
          return tensor_list, pad_sizes
  
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 1 for calling the Shard obj's underlying member method code with example usage: obj._split_tensor(tensor, num_chunks) and returns: shard_list,  pad_sizes, tensor_list,  pad_sizes