code: def forward(self, input: torch.Tensor) -> torch.Tensor:
      if not self.should_gather:
          return super().forward(input)
  
      world_size = self.process_group.size()
      if len(input.shape) == 2 and isinstance(self.linear, FastLinear):
          out_dim = self.linear.weight.shape[0]
  
          if input.shape[0] == 1:
              world_out = input.new_empty(1, out_dim * world_size)
              local_out = input.new_empty(1, out_dim)
              gather_input = local_out
          else:
              world_out = input.new_empty(
                  out_dim * world_size, input.shape[0]
              )
              gather_input = input.new_empty(out_dim, input.shape[0])
              local_out = gather_input.T
  
          torch.mm(input, self.linear.weight.T, out=local_out)
  
          torch.distributed.all_gather_into_tensor(
              world_out, gather_input, group=self.process_group
          )
  
          if input.shape[0] == 1:
              return world_out
          return world_out.T
  
      output = super().forward(input)
      world_output = [
          torch.empty_like(output)
          for _ in range(self.process_group.size())
      ]
      torch.distributed.all_gather(
          world_output, output, group=self.process_group
      )
      world_output = torch.cat(world_output, dim=-1)
      return world_output
  
imports: import os
  import torch
  import torch.distributed
  from torch import nn
  from torch.nn import functional
  from typing import List
  from loguru import logger
  from functools import lru_cache
  from accelerate import init_empty_weights
  from text_generation_server.utils.gptq.quant_linear import QuantLinear
  from text_generation_server.utils.import_utils import IS_CUDA_SYSTEM, IS_ROCM_SYSTEM
  from text_generation_server.utils.log import log_once
label_desc: Use image node_id 3 for calling the TensorParallelHead obj's underlying member method code with example usage: obj.forward(input) and returns: world_output, super, world_out, world_out