code: def forward(
      ctx,
      input_: Tensor,
      input_parallel_mode: ParallelMode,
      weight_parallel_mode: ParallelMode,
      reduce_mean: bool = False,
  ) -> Tensor:
      output = all_reduce(input_, input_parallel_mode)
      output = all_reduce(output, weight_parallel_mode)
      ctx.reduce_mean = reduce_mean
      if reduce_mean:
          reduce_size = gpc.get_world_size(
              input_parallel_mode
          ) * gpc.get_world_size(weight_parallel_mode)
          ctx.reduce_size = reduce_size
          return output.clone() / reduce_size
      return output.clone()
  
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 _ReduceByBatch3D obj's underlying member method code with example usage: obj.forward(ctx, input_, input_parallel_mode, weight_parallel_mode, reduce_mean) and returns: output, unknown