code: def forward(
      ctx,
      input_: Tensor,
      weight: Tensor,
      weight_id: int,
      input_parallel_mode: ParallelMode,
      weight_parallel_mode: ParallelMode,
      output_parallel_mode: ParallelMode,
  ) -> Tensor:
      ctx.weight_id = weight_id
      ctx.input_parallel_mode = input_parallel_mode
      ctx.weight_parallel_mode = weight_parallel_mode
      ctx.output_parallel_mode = output_parallel_mode
  
      input_ = all_gather(input_, 0, input_parallel_mode)
      weight = all_gather(weight, 0, weight_parallel_mode)
      ctx.save_for_backward(input_, weight)
  
      output = torch.matmul(input_, weight)
      output = reduce_scatter(output, 0, output_parallel_mode)
  
      return output
  
imports: from typing import Optional, Tuple
  import torch
  from torch import Tensor
  from torch.cuda.amp import custom_bwd, custom_fwd
  from colossalai.legacy.communication import all_gather, all_reduce, broadcast, reduce, reduce_scatter
  from colossalai.legacy.constants import INPUT_GROUP_3D, WEIGHT_GROUP_3D
  from colossalai.legacy.context.parallel_mode import ParallelMode
  from colossalai.legacy.core import global_context
  from ._utils import get_parallel_mode_from_env, push_async_grad
label_desc: Use image node_id 1 for calling the _Linear3D obj's underlying member method code with example usage: obj.forward(ctx, input_, weight, weight_id, input_parallel_mode, weight_parallel_mode, output_parallel_mode) and returns: output