class_docstr: The DTensor Expand mode. It's replicating the parameters and
  shard the inputs to represent DDP like behavior, it's currently
  a transitent mode before we move to the new data parallel expansion.
class_docstr_tok: ["The","DTensor","Expand","mode",".","It","'s","replicating","the","parameters","and","shard","the","inputs","to","represent","DDP","like","behavior",",","it","'s","currently","a","transitent","mode","before","we","move","to","the","new","data","parallel","expansion","."]
code: def __init__(
      self,
      custom_passes: Optional[
          Callable[[GraphModule], GraphModule]
      ] = None,
  ):
      self._placements_override: Dict[int, List[Placement]] = {}
      if custom_passes is not None:
          self._gm_passes: Callable[
              [GraphModule], GraphModule
          ] = custom_passes
      else:
          # TODO: add a few default passes here.
          self._gm_passes = lambda gm: gm
  
imports: from abc import ABC, abstractmethod
  from typing import Any, Callable, Dict, List, Optional, Tuple
  import torch
  import torch.distributed
  import torch.utils._pytree
  from torch._subclasses import FakeTensorMode
  from torch.distributed._spmd.data_parallel import DataParallelStyle, partition_data_parallel
  from torch.distributed._spmd.distribute import _convert_to_distributed, Schema
  from torch.distributed._tensor import DeviceMesh, Placement, Replicate, Shard
  from torch.fx import GraphModule
label_desc: Use image node_id 1 to create a new DTensorExpandMode object from inherited base classes: ParallelMode with example: obj = DTensorExpandMode(custom_passes)