code: def _create_chunk_sharded_tensor(
      tensor: torch.Tensor,
      rank: int,
      world_size: int,
      num_devices_per_node: int,
      pg: dist.ProcessGroup,
      device: Optional[torch.device] = None,
  ) -> ShardedTensor:
      """
      Shard a tensor to chunks along the first dimension. The local rank will gets its
      corresponding chunk as the local shard to create a ShardedTensor.
      """
      chunks = tensor.chunk(world_size, dim=0)
      if len(chunks) > rank:
          local_shard = chunks[rank].clone()
          offsets = [0 for _ in tensor.size()]
          offsets[0] = math.ceil(tensor.size()[0] / world_size) * rank
          local_shards = [
              Shard.from_tensor_and_offsets(local_shard, offsets, rank)
          ]
      else:
          local_shards = []
  
      # Create a ShardedTensor without invoking communication.
      chunk_sizes = [list(chunk.size()) for chunk in chunks]
      dim0_offsets = [0] + list(
          itertools.accumulate(
              [chunk_size[0] for chunk_size in chunk_sizes]
          )
      )[:-1]
      offsets = [0] * (len(chunk_sizes[0]) - 1)
      chunk_offsets = [[d0] + offsets for d0 in dim0_offsets]
      device_type = (
          distributed_c10d._get_pg_default_device(pg).type
          if device is None
          else device.type
      )
      placements = [
          _get_remote_device_str(r, device_type, num_devices_per_node)
          for r in range(len(chunk_sizes))
      ]
      assert len(chunk_sizes) == len(chunk_offsets) == len(placements)
      shard_metadata = [
          ShardMetadata(offset, size, placement)
          for offset, size, placement in zip(
              chunk_offsets, chunk_sizes, placements
          )
      ]
      sharded_tensor_metadata = ShardedTensorMetadata(
          shards_metadata=shard_metadata,
          size=tensor.size(),
          tensor_properties=TensorProperties(
              dtype=tensor.dtype,
              layout=tensor.layout,
              requires_grad=False,
              memory_format=torch.contiguous_format,
              pin_memory=tensor.is_pinned(),
          ),
      )
      return ShardedTensor._init_from_local_shards_and_global_metadata(
          local_shards,
          sharded_tensor_metadata=sharded_tensor_metadata,
          process_group=pg,
      )
  
file_path: pytorch/torch/distributed/fsdp/_shard_utils.py
imports: import copy
  import itertools
  import math
  from typing import Optional
  import torch
  import torch.distributed
  from torch.distributed import distributed_c10d
  from torch.distributed._shard.sharded_tensor import Shard, ShardedTensor, ShardedTensorMetadata, TensorProperties
  from torch.distributed._shard.sharding_spec import ShardMetadata
  from torch.distributed._tensor import DeviceMesh, DTensor, Replicate, Shard
label_desc: Use image node_id 2 for calling a global function with example usage: _create_chunk_sharded_tensor(tensor, rank, world_size, num_devices_per_node, pg, device) and returns: ShardedTensor