class_docstr: This class is designed for matmul operation with gather forward and reduce-scatter backward.
  
  Args:
      input_ (`torch.Tensor`): input matrix.
      dim (int): the dimension to perform split and gather
      process_group (`torch.distributed.ProcessGroup`): the process group used for collective communication
class_docstr_tok: ["This","class","is","designed","for","matmul","operation","with","gather","forward","and","reduce-scatter","backward",".","Args",":","input_","(","`","torch.Tensor","`",")",":","input","matrix",".","dim","(","int",")",":","the","dimension","to","perform","split","and","gather","process_group","(","`","torch.distributed.ProcessGroup","`",")",":","the","process","group","used","for","collective","communication"]
code: def forward(
      ctx,
      input_,
      weight,
      bias,
      process_group,
      async_grad_reduce_scatter,
      dim,
      overlap,
  ):
      ctx.save_for_backward(input_, weight, bias)
      ctx.use_bias = bias is not None
      ctx.process_group = process_group
      ctx.async_grad_reduce_scatter = async_grad_reduce_scatter
      ctx.dim = dim
      ctx.overlap = overlap
  
      input_parallel = _gather(input_, dim, process_group)
  
      output = torch.matmul(input_parallel, weight)
  
      if bias is not None:
          output = output + bias
      return output
  
imports: import torch
  import torch.distributed
  import torch.nn.functional
label_desc: Use image node_id 1 for calling the _MatmulWithGatherForwardReduceScatterBackward obj's underlying member method code with example usage: obj.forward(ctx, input_, weight, bias, process_group, async_grad_reduce_scatter, dim, overlap) and returns: output