code: def _chunk_tensor(
      tensor: torch.Tensor,
      rank: int,
      world_size: int,
      num_devices_per_node: int,
      pg: dist.ProcessGroup,
  ) -> torch.Tensor:
      if type(tensor) is ShardedTensor:
          assert len(tensor.local_shards()) == 1
  
          inner_param = tensor.local_tensor()
          inner_st = _create_chunk_sharded_tensor(
              inner_param,
              rank,
              world_size,
              num_devices_per_node,
              pg,
          )
  
          outer_local_shard = tensor.local_shards()[0]
          shards: List[Shard] = [
              Shard(inner_st, copy.deepcopy(outer_local_shard.metadata))
          ]
          st_meta = copy.deepcopy(tensor.metadata())
          st_meta.tensor_properties.requires_grad = False
  
          st_outer = (
              ShardedTensor._init_from_local_shards_and_global_metadata(
                  shards,
                  sharded_tensor_metadata=st_meta,
                  process_group=tensor._process_group,
                  init_rrefs=False,
              )
          )
          return st_outer
      elif type(tensor) is DTensor:
          device_mesh = tensor.device_mesh
          assert (
              device_mesh.ndim == 1
          ), "Only 1D DeviceMeshes currently handled"
  
          inner_param = tensor._local_tensor
  
          inner_st = _create_chunk_sharded_tensor(
              inner_param,
              rank,
              world_size,
              torch.cuda.device_count(),
              pg,
          )
  
          dt_pg = _get_dt_pg(tensor)
          # We do this differently here, we create a ST with no local shards then patch it
          shards = [
              Shard(
                  inner_st,
                  _create_shard_md_from_dt(
                      tensor, dist.get_rank(dt_pg)
                  ),
              )
          ]
  
          st_meta = _create_sharded_tensor_md_from_dt(tensor, dt_pg)
          st_meta.tensor_properties.requires_grad = False
  
          st_outer = (
              ShardedTensor._init_from_local_shards_and_global_metadata(
                  shards,
                  sharded_tensor_metadata=st_meta,
                  process_group=dt_pg,
                  init_rrefs=False,
              )
          )
  
          return st_outer
      else:
          return _create_chunk_sharded_tensor(
              tensor,
              rank,
              world_size,
              num_devices_per_node,
              pg,
          )
  
imports: import copy
  from typing import Any, cast, List, Optional, Tuple
  import torch
  import torch.distributed
  import torch.distributed._shard.sharding_spec
  import torch.distributed.distributed_c10d
  from torch.distributed._shard.sharded_tensor import Shard, ShardedTensor, ShardedTensorMetadata, TensorProperties
  from torch.distributed._shard.sharding_spec import ShardMetadata
  from torch.distributed._shard.sharding_spec.chunk_sharding_spec import ChunkShardingSpec
  from torch.distributed._tensor import DeviceMesh, DTensor, Replicate, Shard
  from torch.distributed.device_mesh import _mesh_resources
  from torch.distributed.fsdp._common_utils import _set_fsdp_flattened
  from torch.distributed.fsdp._fsdp_extensions import FSDPExtensions
  from torch.distributed.fsdp._shard_utils import _create_chunk_sharded_tensor
  from torch.distributed.remote_device import _remote_device
  from torch.distributed.tensor.parallel._data_parallel_utils import _flatten_tensor, _unflatten_tensor
label_desc: Use image node_id 8 for calling a global function with example usage: _chunk_tensor(tensor, rank, world_size, num_devices_per_node, pg) and returns: st_outer, st_outer, _create_chunk_sharded_tensor