docstr: We're doing Tensor Parallel with sequence parallelism: we do the matmul and then
  a reduce_scatter of the result.
docstr_tok: ["We","'re","doing","Tensor","Parallel","with","sequence","parallelism",":","we","do","the","matmul","and","then","a","reduce_scatter","of","the","result","."]
code: def forward(self, x):
      """
      We're doing Tensor Parallel with sequence parallelism: we do the matmul and then
      a reduce_scatter of the result.
      """
      out = fused_dense_func(x, self.weight, self.bias)
      reduce_fn = (
          reduce_scatter if self.sequence_parallel else all_reduce
      )
      return reduce_fn(out, self.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 RowParallelLinear obj's underlying member method code with example usage: obj.forward(x) and returns: reduce_fn