class_docstr: Data Parallelism mode.
class_docstr_tok: ["Data","Parallelism","mode","."]
docstr: DataParallel Mode that partition the model and graph to data parallel style
  parallelism (i.e. DDP/FSDP/ZERO-3). It currently supports three different
  parallel styles: "replicate", "fully_shard", and "default". See
  :class:`DataParallelStyle` for more details.
  
  Args:
      parallel_style (str): parallel style to use. Currently supports
          "replicate", "fully_shard", and "default".
  
  Keyword args:
      input_batch_dim (int): the batch dimension of the input tensor.
           default: 0
      custom_passes (Callable[[GraphModule], GraphModule], optional):
          A custom callable that overrides the default graph transformation
          and optimization passes.
docstr_tok: ["DataParallel","Mode","that","partition","the","model","and","graph","to","data","parallel","style","parallelism","(","i.e",".","DDP\/FSDP\/ZERO-3",")",".","It","currently","supports","three","different","parallel","styles",":","``","replicate","''",",","``","fully_shard","''",",","and","``","default","''",".","See",":","class",":","`","DataParallelStyle","`","for","more","details",".","Args",":","parallel_style","(","str",")",":","parallel","style","to","use",".","Currently","supports","``","replicate","''",",","``","fully_shard","''",",","and","``","default","''",".","Keyword","args",":","input_batch_dim","(","int",")",":","the","batch","dimension","of","the","input","tensor",".","default",":","0","custom_passes","(","Callable","[","[","GraphModule","]",",","GraphModule","]",",","optional",")",":","A","custom","callable","that","overrides","the","default","graph","transformation","and","optimization","passes","."]
code: def __init__(
      self,
      parallel_style: str = "replicate",
      *,
      input_batch_dim: int = 0,
      custom_passes: Optional[
          Callable[[GraphModule], GraphModule]
      ] = None,
  ):
      """
      DataParallel Mode that partition the model and graph to data parallel style
      parallelism (i.e. DDP/FSDP/ZERO-3). It currently supports three different
      parallel styles: "replicate", "fully_shard", and "default". See
      :class:`DataParallelStyle` for more details.
  
      Args:
          parallel_style (str): parallel style to use. Currently supports
              "replicate", "fully_shard", and "default".
  
      Keyword args:
          input_batch_dim (int): the batch dimension of the input tensor.
               default: 0
          custom_passes (Callable[[GraphModule], GraphModule], optional):
              A custom callable that overrides the default graph transformation
              and optimization passes.
      """
      if parallel_style == "replicate":
          self.parallel_style = DataParallelStyle.REPLICATE
      elif parallel_style == "fully_shard":
          self.parallel_style = DataParallelStyle.FULLY_SHARD
      elif parallel_style == "default":
          self.parallel_style = DataParallelStyle.DEFAULT
      else:
          raise RuntimeError(
              f"Unknown parallel style: {parallel_style}"
          )
  
      # TODO: what if user passes in a incorrect `input_batch_dim`, how should we
      # detect that and do proper error handling?
      self.input_batch_dim = input_batch_dim
  
      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 DataParallel object from inherited base classes: ParallelMode with example: obj = DataParallel(parallel_style)