code: def _all_to_all_for_tablewise(
      x: torch.Tensor,
      pg: ProcessGroup,
      scatter_strides: List[int],
      gather_strides: List[int],
      forward=True,
  ) -> torch.Tensor:
      world_size = pg.tp_world_size()
      rank = pg.tp_local_rank()
      if world_size == 1:
          return x
      assert (
          x.device.type == "cuda"
      ), f"Currently, the collective function dual_all_to_all only supports nccl backend"
      if forward:
          scatter_list = list(x.split(scatter_strides, 0))
          gather_list = [
              torch.empty(
                  scatter_strides[rank],
                  gather_strides[i],
                  dtype=x.dtype,
                  device=x.device,
              )
              for i in range(world_size)
          ]
          torch.distributed.all_to_all(
              gather_list, scatter_list, group=pg.tp_process_group()
          )
          return torch.cat(gather_list, 1).contiguous()
      else:
          # split on dim 1, lose contiguity
          scatter_list = [
              each.contiguous() for each in x.split(scatter_strides, 1)
          ]
          gather_list = [
              torch.empty(
                  gather_strides[i],
                  scatter_strides[rank],
                  dtype=x.dtype,
                  device=x.device,
              )
              for i in range(world_size)
          ]
          torch.distributed.all_to_all(
              gather_list, scatter_list, group=pg.tp_process_group()
          )
          return torch.cat(gather_list, 0).contiguous()
  
imports: from typing import List, Optional, Union
  import torch
  import torch.distributed
  from colossalai.legacy.global_variables import tensor_parallel_env
  from colossalai.legacy.nn.layer.utils import divide
  from colossalai.legacy.tensor import ColoTensorSpec, ProcessGroup
  from colossalai.tensor import ColoTensor
label_desc: Use image node_id 15 for calling a global function with example usage: _all_to_all_for_tablewise(x, pg, scatter_strides, gather_strides, forward) and returns: x, torch, torch