class_docstr: LayoutConverter is a singleton class which converts the layout of a distributed tensor.
class_docstr_tok: ["LayoutConverter","is","a","singleton","class","which","converts","the","layout","of","a","distributed","tensor","."]
docstr: Get all valid layouts from source_layout with single all-gather operation.
  For the all-gather operation, we just care about the S dimension.
  
  Argument:
      source_layout: the layout to be transformed.
  
  Return:
      valid_spec_dict(Dict[Layout, CommSpec]): all valid layouts from source_layout with single all-gather operation.
  
  Example:
      layout_converter = LayoutConverter()
      physical_mesh_id = torch.arange(0, 4)
      mesh_shape = (2, 2)
      # [[0, 1,
      #  [2, 3]]
      device_mesh = DeviceMesh(physical_mesh_id, mesh_shape, init_process_group=True)
      global_shape = (4, 4, 4)
      dim_partition_dict = {0: [0], 1: [1]}
  
      # [S0,S1,R]
      sharding_spec = ShardingSpec(dim_size=3, dim_partition_dict=dim_partition_dict)
      layout = Layout(device_mesh=device_mesh,
                      sharding_spec=sharding_spec,
                      global_shape=global_shape)
  
      rst_dict = layout_converter.all_gather_transform_layouts(layout)
      for layout, comm_spec in rst_dict.items():
          print(f'{layout.sharding_spec.sharding_sequence}: {comm_spec}')
  
  Output:
      [R, S1, R]: CommSpec:(comm_pattern:GATHER_FWD_SPLIT_BWD, gather_dim:0, shard_dim:0, logical_process_axis:0)
      [S0, R, R]: CommSpec:(comm_pattern:GATHER_FWD_SPLIT_BWD, gather_dim:1, shard_dim:1, logical_process_axis:1)
docstr_tok: ["Get","all","valid","layouts","from","source_layout","with","single","all-gather","operation",".","For","the","all-gather","operation",",","we","just","care","about","the","S","dimension",".","Argument",":","source_layout",":","the","layout","to","be","transformed",".","Return",":","valid_spec_dict","(","Dict","[","Layout",",","CommSpec","]",")",":","all","valid","layouts","from","source_layout","with","single","all-gather","operation",".","Example",":","layout_converter","=","LayoutConverter","(",")","physical_mesh_id","=","torch.arange","(","0",",","4",")","mesh_shape","=","(","2",",","2",")","#","[","[","0",",","1",",","#","[","2",",","3","]","]","device_mesh","=","DeviceMesh","(","physical_mesh_id",",","mesh_shape",",","init_process_group=True",")","global_shape","=","(","4",",","4",",","4",")","dim_partition_dict","=","{","0",":","[","0","]",",","1",":","[","1","]","}","#","[","S0",",","S1",",","R","]","sharding_spec","=","ShardingSpec","(","dim_size=3",",","dim_partition_dict=dim_partition_dict",")","layout","=","Layout","(","device_mesh=device_mesh",",","sharding_spec=sharding_spec",",","global_shape=global_shape",")","rst_dict","=","layout_converter.all_gather_transform_layouts","(","layout",")","for","layout",",","comm_spec","in","rst_dict.items","(",")",":","print","(","f","'","{","layout.sharding_spec.sharding_sequence","}",":","{","comm_spec","}","'",")","Output",":","[","R",",","S1",",","R","]",":","CommSpec",":","(","comm_pattern",":","GATHER_FWD_SPLIT_BWD",",","gather_dim:0",",","shard_dim:0",",","logical_process_axis:0",")","[","S0",",","R",",","R","]",":","CommSpec",":","(","comm_pattern",":","GATHER_FWD_SPLIT_BWD",",","gather_dim:1",",","shard_dim:1",",","logical_process_axis:1",")"]
code: def all_gather_transform_layouts(
      self, source_layout: Layout
  ) -> Dict[Layout, CommSpec]:
      """
      Get all valid layouts from source_layout with single all-gather operation.
      For the all-gather operation, we just care about the S dimension.
  
      Argument:
          source_layout: the layout to be transformed.
  
      Return:
          valid_spec_dict(Dict[Layout, CommSpec]): all valid layouts from source_layout with single all-gather operation.
  
      Example:
          layout_converter = LayoutConverter()
          physical_mesh_id = torch.arange(0, 4)
          mesh_shape = (2, 2)
          # [[0, 1,
          #  [2, 3]]
          device_mesh = DeviceMesh(physical_mesh_id, mesh_shape, init_process_group=True)
          global_shape = (4, 4, 4)
          dim_partition_dict = {0: [0], 1: [1]}
  
          # [S0,S1,R]
          sharding_spec = ShardingSpec(dim_size=3, dim_partition_dict=dim_partition_dict)
          layout = Layout(device_mesh=device_mesh,
                          sharding_spec=sharding_spec,
                          global_shape=global_shape)
  
          rst_dict = layout_converter.all_gather_transform_layouts(layout)
          for layout, comm_spec in rst_dict.items():
              print(f'{layout.sharding_spec.sharding_sequence}: {comm_spec}')
  
      Output:
          [R, S1, R]: CommSpec:(comm_pattern:GATHER_FWD_SPLIT_BWD, gather_dim:0, shard_dim:0, logical_process_axis:0)
          [S0, R, R]: CommSpec:(comm_pattern:GATHER_FWD_SPLIT_BWD, gather_dim:1, shard_dim:1, logical_process_axis:1)
      """
      valid_spec_dict = {}
      comm_pattern = CollectiveCommPattern.GATHER_FWD_SPLIT_BWD
      source_spec = source_layout.sharding_spec
  
      # the key of the dict is the axis
      # the value is the process group
      current_rank = (
          source_layout.device_mesh._global_rank_of_current_process
      )
      process_group_dict = (
          source_layout.device_mesh._process_group_dict[current_rank]
      )
  
      for target_pair in source_spec.dim_partition_dict.items():
          shard_list = all_gather_simulator(target_pair)
          index = target_pair[0]
          new_dim_partition_dict = deepcopy(
              source_spec.dim_partition_dict
          )
  
          # We won't add empty list into dim_partition_dict
          # The key will be popped if the related shard_list is empty
          if shard_list:
              new_dim_partition_dict[index] = shard_list
          else:
              new_dim_partition_dict.pop(index)
  
          # generate the CommSpec to record the action of source_sharding_spec->new_sharding_spec
          gather_dim = index
          logical_process_axis = target_pair[1][-1]
          comm_spec = CommSpec(
              comm_pattern,
              process_group_dict=process_group_dict,
              gather_dim=gather_dim,
              # shard_dim will be used during backward
              shard_dim=gather_dim,
              logical_process_axis=logical_process_axis,
          )
  
          # generate new sharding spec
          try:
              new_sharding_spec = ShardingSpec(
                  source_spec.dims,
                  dim_partition_dict=new_dim_partition_dict,
              )
              new_layout = Layout(
                  device_mesh=source_layout.device_mesh,
                  sharding_spec=new_sharding_spec,
                  global_shape=source_layout.global_shape,
              )
  
              valid_spec_dict[new_layout] = comm_spec
          except LayoutException:
              pass
      return valid_spec_dict
  
imports: import math
  from copy import deepcopy
  from dataclasses import dataclass
  from typing import Dict, List, Tuple
  import torch
  import torch.distributed
  from colossalai.context.singleton_meta import SingletonMeta
  from colossalai.tensor.d_tensor.comm_spec import *
  from colossalai.tensor.d_tensor.layout import Layout
  from colossalai.tensor.d_tensor.misc import LayoutException
  from colossalai.tensor.utils import all_gather_simulator, all_to_all_simulator, shard_simulator
  from .sharding_spec import ShardingSpec
  from .utils import get_comm_cost
label_desc: Use image node_id 6 for calling the LayoutConverter obj's underlying member method code with example usage: obj.all_gather_transform_layouts(source_layout) and returns: valid_spec_dict