class_docstr: DTensorExtension is the TensorFlattener extension needed for 2D FSDP + TP.
  
  This is the implementation for FSDPExtensions defined in
  https://github.com/pytorch/pytorch/blob/main/torch/distributed/fsdp/_fsdp_extensions.py
class_docstr_tok: ["DTensorExtension","is","the","TensorFlattener","extension","needed","for","2D","FSDP","+","TP",".","This","is","the","implementation","for","FSDPExtensions","defined","in","https",":","\/\/github.com\/pytorch\/pytorch\/blob\/main\/torch\/distributed\/fsdp\/_fsdp_extensions.py"]
code: def chunk_dtensor(
      self,
      tensor: torch.Tensor,
      rank: int,
      device_mesh: DeviceMesh,
  ) -> torch.Tensor:
      return _chunk_dtensor(tensor, rank, device_mesh)
  
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 5 for calling the DTensorExtensions obj's underlying member method code with example usage: obj.chunk_dtensor(tensor, rank, device_mesh) and returns: _chunk_dtensor