docstr: If process_group is not None, we're doing Tensor Parallel with sequence parallelism:
  we do an all_gather of x before doing the matmul.
docstr_tok: ["If","process_group","is","not","None",",","we","'re","doing","Tensor","Parallel","with","sequence","parallelism",":","we","do","an","all_gather","of","x","before","doing","the","matmul","."]
code: def forward(self, x, process_group=None):
      """
      If process_group is not None, we're doing Tensor Parallel with sequence parallelism:
      we do an all_gather of x before doing the matmul.
      """
      return fused_dense_func(
          x,
          self.weight,
          self.bias,
          return_residual=self.return_residual,
          process_group=process_group,
      )
  
imports: from functools import partial
  from typing import Optional
  import fused_dense_lib
  import torch
  import torch.nn
  import torch.nn.functional
  from torch import Tensor
  from torch.cuda.amp import custom_bwd, custom_fwd
  from torch.distributed import ProcessGroup
  from flash_attn.ops.activations import gelu_bwd, relu_bwd, sqrelu_bwd, sqrelu_fwd
  from flash_attn.utils.distributed import all_gather_raw, all_reduce, all_reduce_raw, reduce_scatter, reduce_scatter_raw
label_desc: Use image node_id 2 for calling the FusedDense obj's underlying member method code with example usage: obj.forward(x, process_group) and returns: fused_dense_func