class_docstr: WhereGenerator is a generic class to generate strategies for Where operation.
class_docstr_tok: ["WhereGenerator","is","a","generic","class","to","generate","strategies","for","Where","operation","."]
code: def _generate_strategy_with_dim_partition(self, dim_partition):
      dim_partition_dict_mapping = {
          "condition": dim_partition,
          "x": dim_partition,
          "y": dim_partition,
          "output": dim_partition,
      }
  
      sharding_spec_mapping = self.to_sharding_spec_mapping(
          dim_partition_dict_mapping
      )
  
      name = f'{sharding_spec_mapping["output"].sharding_sequence} = {sharding_spec_mapping["condition"].sharding_sequence} x {sharding_spec_mapping["x"].sharding_sequence} x {sharding_spec_mapping["y"].sharding_sequence}'
      communication_action_mapping = {}
  
      strategy = self.get_sharding_strategy(
          name=name,
          sharding_spec_mapping=sharding_spec_mapping,
          communication_action_mapping=communication_action_mapping,
      )
  
      return strategy
  
file_path: ColossalAI/colossalai/auto_parallel/tensor_shard/node_handler/strategy/where_generator.py
imports: import copy
  from typing import List
  from colossalai.auto_parallel.tensor_shard.sharding_strategy import MemoryCost, ShardingStrategy, TrainCycleItem
  from colossalai.auto_parallel.tensor_shard.utils import enumerate_all_possible_1d_sharding, enumerate_all_possible_2d_sharding, ignore_sharding_exception
  from .strategy_generator import StrategyGenerator
label_desc: Use image node_id 4 for calling the WhereGenerator obj's underlying member method code with example usage: obj._generate_strategy_with_dim_partition(dim_partition) and returns: strategy