We are currently solving the following issue within our repository. Here is the issue text:
--- BEGIN ISSUE ---
[tensor] fix some unittests

[tensor] fix some unittests

[tensor] fix some unittests


--- END ISSUE ---

Below are some code segments, each from a relevant file. One or more of these files may contain bugs.
--- BEGIN FILES ---
Path: `colossalai/auto_parallel/tensor_shard/deprecated/op_handler/where_handler.py`
Content:
```
1 import operator
2 import warnings
3 from copy import deepcopy
4 from functools import reduce
5 from typing import Dict, List
6 
7 import torch
8 
9 from colossalai.auto_parallel.tensor_shard.deprecated._utils import (enumerate_all_possible_1d_sharding,
10                                                                      enumerate_all_possible_2d_sharding,
11                                                                      ignore_sharding_exception)
12 from colossalai.auto_parallel.tensor_shard.deprecated.sharding_strategy import (ShardingStrategy, StrategiesVector)
13 from colossalai.tensor.shape_consistency import ShapeConsistencyManager
14 from colossalai.tensor.sharding_spec import ShardingSpec
15 
16 from .operator_handler import OperatorHandler
17 
18 __all__ = ['WhereHandler']
19 
20 
21 class WhereHandler(OperatorHandler):
22     """
23     An OperatorHandler which deals with the sharding strategies of torch.where.
24     """
25 
26     def __init__(self, *args, **kwargs):
27         # TODO: x or y could be scalar
28         super().__init__(*args, **kwargs)
29         assert len(self.predecessor_node) == 3
30         self.condition_data = self.predecessor_node[0]._meta_data
31         self.x_data = self.predecessor_node[1]._meta_data
32         self.y_data = self.predecessor_node[2]._meta_data
33         self.condition = self.predecessor_node[0]
34         self.x = self.predecessor_node[1]
35         self.y = self.predecessor_node[2]
36         self.output_data = self.node._meta_data
37 
38     def _generate_sharding_spec(self, input_: torch.Tensor, dim_partition_dict: Dict[int, List[int]]) -> ShardingSpec:
39         shape = list(input_.shape)
40 
41         # padding the shape to the same length as output_data
42         while len(shape) < self.output_data.dim():
43             shape.insert(0, 1)
44         shape = torch.Size(shape)
45 
46         # if the sharding happens on a size one dimension, we should record it as R.
47         processed_dim_partition_dict = deepcopy(dim_partition_dict)
48         for dim_index, _ in dim_partition_dict.items():
49             if shape[dim_index] == 1:
50                 processed_dim_partition_dict.pop(dim_index)
51         for dim_index, sharding_index_list in processed_dim_partition_dict.items():
52             sharding_list = [self.device_mesh.mesh_shape[sharding_index] for sharding_index in sharding_index_list]
53             sharding_size = reduce(operator.mul, sharding_list, 1)
54             assert shape[
55                 dim_index] % sharding_size == 0, f'we cannot shard the {dim_index} dimension of tensor into {sharding_size} partitions.'
56         sharding_spec = ShardingSpec(device_mesh=self.device_mesh,
57                                      entire_shape=shape,
58                                      dim_partition_dict=processed_dim_partition_dict)
59 
60         return sharding_spec
61 
62     def _generate_compute_cost(self, total_sharding_size):
63         lhs_matrix_shape = self.lhs_data.shape[-2:]
64         rhs_matrix_shape = self.rhs_data.shape[-2:]
65         batch_dimensions_shape = self.output_data.shape[:-2]
66         batch_dimensions_product = reduce(operator.mul, batch_dimensions_shape, 1)
67         compute_cost = reduce(
68             operator.mul, lhs_matrix_shape) * rhs_matrix_shape[0] * batch_dimensions_product * 2 / total_sharding_size
69         return compute_cost
70 
71     def _generate_resharding_costs(self, sharding_specs):
72         # The resharding_cost of weight is counted due to sharing weight cases.
73         dtype = self.node._meta_data.dtype
74         nodes = self.predecessor_node
75         resharding_costs = {}
76         size_per_elem_bytes = torch.tensor([], dtype=dtype).element_size()
77 
78         # shape consistency manager is a singleton class
79         shape_consistency_manager = ShapeConsistencyManager()
80 
81         for input_node, input_spec in zip(nodes, sharding_specs):
82             resharding_costs[input_node] = []
83             for strategy in input_node.strategies_vector:
84                 input_sharding_spec = strategy.output_sharding_spec
85                 assert isinstance(input_sharding_spec, ShardingSpec), f'The input node should NOT be a tuple of tensor.'
86                 # if the input shape is smaller than the target input, we will fill the input to the same length as target.
87                 # Then, use the padded input sharding spec to compute the resharding cost.
88                 if len(input_sharding_spec.entire_shape) < len(input_spec.entire_shape):
89                     new_entire_shape = list(input_sharding_spec.entire_shape)
90                     while len(new_entire_shape) < len(input_spec.entire_shape):
91                         new_entire_shape.insert(0, 1)
92                     new_entire_shape = torch.Size(new_entire_shape)
93                     new_device_mesh = input_sharding_spec.device_mesh
94                     new_dim_partition_dict = input_sharding_spec.dim_partition_dict
95                     input_sharding_spec = ShardingSpec(device_mesh=new_device_mesh,
96                                                        entire_shape=new_entire_shape,
97                                                        dim_partition_dict=new_dim_partition_dict)
98 
99                 # compute the resharding cost
100                 _, _, total_resharding_cost = shape_consistency_manager.shape_consistency(
101                     input_sharding_spec, input_spec)
102                 total_resharding_cost = total_resharding_cost['total']
103                 # we need multiply the size of elem dtype to get correct communication cost
104                 resharding_cost = total_resharding_cost * size_per_elem_bytes
105                 resharding_costs[input_node].append(resharding_cost)
106 
107         return resharding_costs
108 
109     def _convert_partition_dict_to_sharding_spec(self, dim_partition_list):
110 
111         sharding_spec_list = []
112         check_duplicated_list = []
113         for output_dim_partition_dict in dim_partition_list:
114             try:
115                 output_sharding_spec = self._generate_sharding_spec(self.output_data, output_dim_partition_dict)
116             except AssertionError as e:
117                 warnings.warn(f'{e}')
118                 break
119             sharding_seq = output_sharding_spec.sharding_sequence
120             if sharding_seq not in check_duplicated_list:
121                 check_duplicated_list.append(sharding_seq)
122                 sharding_spec_list.append(output_sharding_spec)
123 
124         return sharding_spec_list
125 
126     def _enumerate_all_possible_output(self, mesh_dim_0, mesh_dim_1):
127         # use mesh_dim_0, mesh_dim_1 instead of constant 0, 1 in here for N-D device mesh scaliablity.
128 
129         output_dim_partition_list = []
130         dim_size = self.output_data.dim()
131         # enumerate all the 2D sharding cases
132         sharding_list_2d = enumerate_all_possible_2d_sharding(mesh_dim_0, mesh_dim_1, dim_size)
133         output_dim_partition_list.extend(sharding_list_2d)
134 
135         # enumerate all the 1D sharding cases
136         sharding_list_1d_on_dim_0 = enumerate_all_possible_1d_sharding(mesh_dim_0, dim_size)
137         output_dim_partition_list.extend(sharding_list_1d_on_dim_0)
138         sharding_list_1d_on_dim_1 = enumerate_all_possible_1d_sharding(mesh_dim_1, dim_size)
139         output_dim_partition_list.extend(sharding_list_1d_on_dim_1)
140 
141         # add empty dict for fully replicated case
142         output_dim_partition_list.append({})
143         output_sharding_spec_list = self._convert_partition_dict_to_sharding_spec(output_dim_partition_list)
144 
145         return output_sharding_spec_list
146 
147     @ignore_sharding_exception
148     def _register_strategy(self, output_sharding_spec):
149         dim_partition_dict_for_input = output_sharding_spec.dim_partition_dict
150         sharding_spec_for_condition = self._generate_sharding_spec(self.condition_data, dim_partition_dict_for_input)
151         sharding_spec_for_x = self._generate_sharding_spec(self.x_data, dim_partition_dict_for_input)
152         sharding_spec_for_y = self._generate_sharding_spec(self.y_data, dim_partition_dict_for_input)
153 
154         name = f'{output_sharding_spec.sharding_sequence} = {sharding_spec_for_condition.sharding_sequence} x {sharding_spec_for_x.sharding_sequence} x {sharding_spec_for_y.sharding_sequence}'
155         dim_partition_dict_for_output = output_sharding_spec.dim_partition_dict
156 
157         # generate resharding cost for this strategy
158         resharding_costs = self._generate_resharding_costs(
159             [sharding_spec_for_condition, sharding_spec_for_x, sharding_spec_for_y])
160 
161         # compute the computation cost of this strategy
162         sharding_dims = []
163         for mesh_dims in dim_partition_dict_for_output.values():
164             for mesh_dim in mesh_dims:
165                 sharding_dims.append(self.device_mesh.shape[mesh_dim])
166         sharding_size = reduce(operator.mul, sharding_dims, 1)
167         memory_cost = self.output_data.numel() / sharding_size
168         compute_cost = memory_cost
169         communication_cost = 0
170 
171         sharding_strategies = ShardingStrategy(name,
172                                                output_sharding_spec=output_sharding_spec,
173                                                compute_cost=compute_cost,
174                                                communication_cost=communication_cost,
175                                                memory_cost=memory_cost,
176                                                resharding_costs=resharding_costs,
177                                                input_shardings=(sharding_spec_for_condition, sharding_spec_for_x,
178                                                                 sharding_spec_for_y))
179 
180         self.strategies_vector.append(sharding_strategies)
181 
182     def register_strategy(self) -> StrategiesVector:
183         MESH_DIM_LIST = [0, 1]
184         output_sharding_specs = self._enumerate_all_possible_output(MESH_DIM_LIST[0], MESH_DIM_LIST[1])
185         for output_sharding_spec in output_sharding_specs:
186             self._register_strategy(output_sharding_spec)
187
```
--- END FILES ---

Please first localize the bug based on the issue statement, and then generate a patch according to the `git diff` format fenced by three backticks.

Here is an example:

```diff
diff --git a/examples/server_async.py b/examples/server_async.py
--- a/examples/server_async.py
+++ b/examples/server_async.py
@@ -313,4 +313,4 @@
 
 
 if __name__ == "__main__":
-    asyncio.run(run_async_server("."), debug=True)
+    asyncio.run(run_async_server(), debug=True)
diff --git a/examples/server_sync.py b/examples/server_sync.py
--- a/examples/server_sync.py
+++ b/examples/server_sync.py
@@ -313,5 +313,5 @@
 
 
 if __name__ == "__main__":
-    server = run_sync_server(".")
+    server = run_sync_server()
     server.shutdown()

```

diff --git a/colossalai/auto_parallel/tensor_shard/deprecated/op_handler/where_handler.py b/colossalai/auto_parallel/tensor_shard/deprecated/op_handler/where_handler.py
--- a/colossalai/auto_parallel/tensor_shard/deprecated/op_handler/where_handler.py
+++ b/colossalai/auto_parallel/tensor_shard/deprecated/op_handler/where_handler.py
@@ -6,10 +6,12 @@
 
 import torch
 
-from colossalai.auto_parallel.tensor_shard.deprecated._utils import (enumerate_all_possible_1d_sharding,
-                                                                     enumerate_all_possible_2d_sharding,
-                                                                     ignore_sharding_exception)
-from colossalai.auto_parallel.tensor_shard.deprecated.sharding_strategy import (ShardingStrategy, StrategiesVector)
+from colossalai.auto_parallel.tensor_shard.deprecated._utils import (
+    enumerate_all_possible_1d_sharding,
+    enumerate_all_possible_2d_sharding,
+    ignore_sharding_exception,
+)
+from colossalai.auto_parallel.tensor_shard.deprecated.sharding_strategy import ShardingStrategy, StrategiesVector
 from colossalai.tensor.shape_consistency import ShapeConsistencyManager
 from colossalai.tensor.sharding_spec import ShardingSpec

{"golden_diff": "diff --git a/colossalai/auto_parallel/tensor_shard/deprecated/op_handler/where_handler.py b/colossalai/auto_parallel/tensor_shard/deprecated/op_handler/where_handler.py\n--- a/colossalai/auto_parallel/tensor_shard/deprecated/op_handler/where_handler.py\n+++ b/colossalai/auto_parallel/tensor_shard/deprecated/op_handler/where_handler.py\n@@ -6,10 +6,12 @@\n \n import torch\n \n-from colossalai.auto_parallel.tensor_shard.deprecated._utils import (enumerate_all_possible_1d_sharding,\n-                                                                     enumerate_all_possible_2d_sharding,\n-                                                                     ignore_sharding_exception)\n-from colossalai.auto_parallel.tensor_shard.deprecated.sharding_strategy import (ShardingStrategy, StrategiesVector)\n+from colossalai.auto_parallel.tensor_shard.deprecated._utils import (\n+    enumerate_all_possible_1d_sharding,\n+    enumerate_all_possible_2d_sharding,\n+    ignore_sharding_exception,\n+)\n+from colossalai.auto_parallel.tensor_shard.deprecated.sharding_strategy import ShardingStrategy, StrategiesVector\n from colossalai.tensor.shape_consistency import ShapeConsistencyManager\n from colossalai.tensor.sharding_spec import ShardingSpec\n", "issue": "[tensor] fix some unittests\n\n[tensor] fix some unittests\n\n[tensor] fix some unittests\n\n", "before_files": [{"content": "import operator\nimport warnings\nfrom copy import deepcopy\nfrom functools import reduce\nfrom typing import Dict, List\n\nimport torch\n\nfrom colossalai.auto_parallel.tensor_shard.deprecated._utils import (enumerate_all_possible_1d_sharding,\n                                                                     enumerate_all_possible_2d_sharding,\n                                                                     ignore_sharding_exception)\nfrom colossalai.auto_parallel.tensor_shard.deprecated.sharding_strategy import (ShardingStrategy, StrategiesVector)\nfrom colossalai.tensor.shape_consistency import ShapeConsistencyManager\nfrom colossalai.tensor.sharding_spec import ShardingSpec\n\nfrom .operator_handler import OperatorHandler\n\n__all__ = ['WhereHandler']\n\n\nclass WhereHandler(OperatorHandler):\n    \"\"\"\n    An OperatorHandler which deals with the sharding strategies of torch.where.\n    \"\"\"\n\n    def __init__(self, *args, **kwargs):\n        # TODO: x or y could be scalar\n        super().__init__(*args, **kwargs)\n        assert len(self.predecessor_node) == 3\n        self.condition_data = self.predecessor_node[0]._meta_data\n        self.x_data = self.predecessor_node[1]._meta_data\n        self.y_data = self.predecessor_node[2]._meta_data\n        self.condition = self.predecessor_node[0]\n        self.x = self.predecessor_node[1]\n        self.y = self.predecessor_node[2]\n        self.output_data = self.node._meta_data\n\n    def _generate_sharding_spec(self, input_: torch.Tensor, dim_partition_dict: Dict[int, List[int]]) -> ShardingSpec:\n        shape = list(input_.shape)\n\n        # padding the shape to the same length as output_data\n        while len(shape) < self.output_data.dim():\n            shape.insert(0, 1)\n        shape = torch.Size(shape)\n\n        # if the sharding happens on a size one dimension, we should record it as R.\n        processed_dim_partition_dict = deepcopy(dim_partition_dict)\n        for dim_index, _ in dim_partition_dict.items():\n            if shape[dim_index] == 1:\n                processed_dim_partition_dict.pop(dim_index)\n        for dim_index, sharding_index_list in processed_dim_partition_dict.items():\n            sharding_list = [self.device_mesh.mesh_shape[sharding_index] for sharding_index in sharding_index_list]\n            sharding_size = reduce(operator.mul, sharding_list, 1)\n            assert shape[\n                dim_index] % sharding_size == 0, f'we cannot shard the {dim_index} dimension of tensor into {sharding_size} partitions.'\n        sharding_spec = ShardingSpec(device_mesh=self.device_mesh,\n                                     entire_shape=shape,\n                                     dim_partition_dict=processed_dim_partition_dict)\n\n        return sharding_spec\n\n    def _generate_compute_cost(self, total_sharding_size):\n        lhs_matrix_shape = self.lhs_data.shape[-2:]\n        rhs_matrix_shape = self.rhs_data.shape[-2:]\n        batch_dimensions_shape = self.output_data.shape[:-2]\n        batch_dimensions_product = reduce(operator.mul, batch_dimensions_shape, 1)\n        compute_cost = reduce(\n            operator.mul, lhs_matrix_shape) * rhs_matrix_shape[0] * batch_dimensions_product * 2 / total_sharding_size\n        return compute_cost\n\n    def _generate_resharding_costs(self, sharding_specs):\n        # The resharding_cost of weight is counted due to sharing weight cases.\n        dtype = self.node._meta_data.dtype\n        nodes = self.predecessor_node\n        resharding_costs = {}\n        size_per_elem_bytes = torch.tensor([], dtype=dtype).element_size()\n\n        # shape consistency manager is a singleton class\n        shape_consistency_manager = ShapeConsistencyManager()\n\n        for input_node, input_spec in zip(nodes, sharding_specs):\n            resharding_costs[input_node] = []\n            for strategy in input_node.strategies_vector:\n                input_sharding_spec = strategy.output_sharding_spec\n                assert isinstance(input_sharding_spec, ShardingSpec), f'The input node should NOT be a tuple of tensor.'\n                # if the input shape is smaller than the target input, we will fill the input to the same length as target.\n                # Then, use the padded input sharding spec to compute the resharding cost.\n                if len(input_sharding_spec.entire_shape) < len(input_spec.entire_shape):\n                    new_entire_shape = list(input_sharding_spec.entire_shape)\n                    while len(new_entire_shape) < len(input_spec.entire_shape):\n                        new_entire_shape.insert(0, 1)\n                    new_entire_shape = torch.Size(new_entire_shape)\n                    new_device_mesh = input_sharding_spec.device_mesh\n                    new_dim_partition_dict = input_sharding_spec.dim_partition_dict\n                    input_sharding_spec = ShardingSpec(device_mesh=new_device_mesh,\n                                                       entire_shape=new_entire_shape,\n                                                       dim_partition_dict=new_dim_partition_dict)\n\n                # compute the resharding cost\n                _, _, total_resharding_cost = shape_consistency_manager.shape_consistency(\n                    input_sharding_spec, input_spec)\n                total_resharding_cost = total_resharding_cost['total']\n                # we need multiply the size of elem dtype to get correct communication cost\n                resharding_cost = total_resharding_cost * size_per_elem_bytes\n                resharding_costs[input_node].append(resharding_cost)\n\n        return resharding_costs\n\n    def _convert_partition_dict_to_sharding_spec(self, dim_partition_list):\n\n        sharding_spec_list = []\n        check_duplicated_list = []\n        for output_dim_partition_dict in dim_partition_list:\n            try:\n                output_sharding_spec = self._generate_sharding_spec(self.output_data, output_dim_partition_dict)\n            except AssertionError as e:\n                warnings.warn(f'{e}')\n                break\n            sharding_seq = output_sharding_spec.sharding_sequence\n            if sharding_seq not in check_duplicated_list:\n                check_duplicated_list.append(sharding_seq)\n                sharding_spec_list.append(output_sharding_spec)\n\n        return sharding_spec_list\n\n    def _enumerate_all_possible_output(self, mesh_dim_0, mesh_dim_1):\n        # use mesh_dim_0, mesh_dim_1 instead of constant 0, 1 in here for N-D device mesh scaliablity.\n\n        output_dim_partition_list = []\n        dim_size = self.output_data.dim()\n        # enumerate all the 2D sharding cases\n        sharding_list_2d = enumerate_all_possible_2d_sharding(mesh_dim_0, mesh_dim_1, dim_size)\n        output_dim_partition_list.extend(sharding_list_2d)\n\n        # enumerate all the 1D sharding cases\n        sharding_list_1d_on_dim_0 = enumerate_all_possible_1d_sharding(mesh_dim_0, dim_size)\n        output_dim_partition_list.extend(sharding_list_1d_on_dim_0)\n        sharding_list_1d_on_dim_1 = enumerate_all_possible_1d_sharding(mesh_dim_1, dim_size)\n        output_dim_partition_list.extend(sharding_list_1d_on_dim_1)\n\n        # add empty dict for fully replicated case\n        output_dim_partition_list.append({})\n        output_sharding_spec_list = self._convert_partition_dict_to_sharding_spec(output_dim_partition_list)\n\n        return output_sharding_spec_list\n\n    @ignore_sharding_exception\n    def _register_strategy(self, output_sharding_spec):\n        dim_partition_dict_for_input = output_sharding_spec.dim_partition_dict\n        sharding_spec_for_condition = self._generate_sharding_spec(self.condition_data, dim_partition_dict_for_input)\n        sharding_spec_for_x = self._generate_sharding_spec(self.x_data, dim_partition_dict_for_input)\n        sharding_spec_for_y = self._generate_sharding_spec(self.y_data, dim_partition_dict_for_input)\n\n        name = f'{output_sharding_spec.sharding_sequence} = {sharding_spec_for_condition.sharding_sequence} x {sharding_spec_for_x.sharding_sequence} x {sharding_spec_for_y.sharding_sequence}'\n        dim_partition_dict_for_output = output_sharding_spec.dim_partition_dict\n\n        # generate resharding cost for this strategy\n        resharding_costs = self._generate_resharding_costs(\n            [sharding_spec_for_condition, sharding_spec_for_x, sharding_spec_for_y])\n\n        # compute the computation cost of this strategy\n        sharding_dims = []\n        for mesh_dims in dim_partition_dict_for_output.values():\n            for mesh_dim in mesh_dims:\n                sharding_dims.append(self.device_mesh.shape[mesh_dim])\n        sharding_size = reduce(operator.mul, sharding_dims, 1)\n        memory_cost = self.output_data.numel() / sharding_size\n        compute_cost = memory_cost\n        communication_cost = 0\n\n        sharding_strategies = ShardingStrategy(name,\n                                               output_sharding_spec=output_sharding_spec,\n                                               compute_cost=compute_cost,\n                                               communication_cost=communication_cost,\n                                               memory_cost=memory_cost,\n                                               resharding_costs=resharding_costs,\n                                               input_shardings=(sharding_spec_for_condition, sharding_spec_for_x,\n                                                                sharding_spec_for_y))\n\n        self.strategies_vector.append(sharding_strategies)\n\n    def register_strategy(self) -> StrategiesVector:\n        MESH_DIM_LIST = [0, 1]\n        output_sharding_specs = self._enumerate_all_possible_output(MESH_DIM_LIST[0], MESH_DIM_LIST[1])\n        for output_sharding_spec in output_sharding_specs:\n            self._register_strategy(output_sharding_spec)\n", "path": "colossalai/auto_parallel/tensor_shard/deprecated/op_handler/where_handler.py"}], "after_files": [{"content": "import operator\nimport warnings\nfrom copy import deepcopy\nfrom functools import reduce\nfrom typing import Dict, List\n\nimport torch\n\nfrom colossalai.auto_parallel.tensor_shard.deprecated._utils import (\n    enumerate_all_possible_1d_sharding,\n    enumerate_all_possible_2d_sharding,\n    ignore_sharding_exception,\n)\nfrom colossalai.auto_parallel.tensor_shard.deprecated.sharding_strategy import ShardingStrategy, StrategiesVector\nfrom colossalai.tensor.shape_consistency import ShapeConsistencyManager\nfrom colossalai.tensor.sharding_spec import ShardingSpec\n\nfrom .operator_handler import OperatorHandler\n\n__all__ = ['WhereHandler']\n\n\nclass WhereHandler(OperatorHandler):\n    \"\"\"\n    An OperatorHandler which deals with the sharding strategies of torch.where.\n    \"\"\"\n\n    def __init__(self, *args, **kwargs):\n        # TODO: x or y could be scalar\n        super().__init__(*args, **kwargs)\n        assert len(self.predecessor_node) == 3\n        self.condition_data = self.predecessor_node[0]._meta_data\n        self.x_data = self.predecessor_node[1]._meta_data\n        self.y_data = self.predecessor_node[2]._meta_data\n        self.condition = self.predecessor_node[0]\n        self.x = self.predecessor_node[1]\n        self.y = self.predecessor_node[2]\n        self.output_data = self.node._meta_data\n\n    def _generate_sharding_spec(self, input_: torch.Tensor, dim_partition_dict: Dict[int, List[int]]) -> ShardingSpec:\n        shape = list(input_.shape)\n\n        # padding the shape to the same length as output_data\n        while len(shape) < self.output_data.dim():\n            shape.insert(0, 1)\n        shape = torch.Size(shape)\n\n        # if the sharding happens on a size one dimension, we should record it as R.\n        processed_dim_partition_dict = deepcopy(dim_partition_dict)\n        for dim_index, _ in dim_partition_dict.items():\n            if shape[dim_index] == 1:\n                processed_dim_partition_dict.pop(dim_index)\n        for dim_index, sharding_index_list in processed_dim_partition_dict.items():\n            sharding_list = [self.device_mesh.mesh_shape[sharding_index] for sharding_index in sharding_index_list]\n            sharding_size = reduce(operator.mul, sharding_list, 1)\n            assert shape[\n                dim_index] % sharding_size == 0, f'we cannot shard the {dim_index} dimension of tensor into {sharding_size} partitions.'\n        sharding_spec = ShardingSpec(device_mesh=self.device_mesh,\n                                     entire_shape=shape,\n                                     dim_partition_dict=processed_dim_partition_dict)\n\n        return sharding_spec\n\n    def _generate_compute_cost(self, total_sharding_size):\n        lhs_matrix_shape = self.lhs_data.shape[-2:]\n        rhs_matrix_shape = self.rhs_data.shape[-2:]\n        batch_dimensions_shape = self.output_data.shape[:-2]\n        batch_dimensions_product = reduce(operator.mul, batch_dimensions_shape, 1)\n        compute_cost = reduce(\n            operator.mul, lhs_matrix_shape) * rhs_matrix_shape[0] * batch_dimensions_product * 2 / total_sharding_size\n        return compute_cost\n\n    def _generate_resharding_costs(self, sharding_specs):\n        # The resharding_cost of weight is counted due to sharing weight cases.\n        dtype = self.node._meta_data.dtype\n        nodes = self.predecessor_node\n        resharding_costs = {}\n        size_per_elem_bytes = torch.tensor([], dtype=dtype).element_size()\n\n        # shape consistency manager is a singleton class\n        shape_consistency_manager = ShapeConsistencyManager()\n\n        for input_node, input_spec in zip(nodes, sharding_specs):\n            resharding_costs[input_node] = []\n            for strategy in input_node.strategies_vector:\n                input_sharding_spec = strategy.output_sharding_spec\n                assert isinstance(input_sharding_spec, ShardingSpec), f'The input node should NOT be a tuple of tensor.'\n                # if the input shape is smaller than the target input, we will fill the input to the same length as target.\n                # Then, use the padded input sharding spec to compute the resharding cost.\n                if len(input_sharding_spec.entire_shape) < len(input_spec.entire_shape):\n                    new_entire_shape = list(input_sharding_spec.entire_shape)\n                    while len(new_entire_shape) < len(input_spec.entire_shape):\n                        new_entire_shape.insert(0, 1)\n                    new_entire_shape = torch.Size(new_entire_shape)\n                    new_device_mesh = input_sharding_spec.device_mesh\n                    new_dim_partition_dict = input_sharding_spec.dim_partition_dict\n                    input_sharding_spec = ShardingSpec(device_mesh=new_device_mesh,\n                                                       entire_shape=new_entire_shape,\n                                                       dim_partition_dict=new_dim_partition_dict)\n\n                # compute the resharding cost\n                _, _, total_resharding_cost = shape_consistency_manager.shape_consistency(\n                    input_sharding_spec, input_spec)\n                total_resharding_cost = total_resharding_cost['total']\n                # we need multiply the size of elem dtype to get correct communication cost\n                resharding_cost = total_resharding_cost * size_per_elem_bytes\n                resharding_costs[input_node].append(resharding_cost)\n\n        return resharding_costs\n\n    def _convert_partition_dict_to_sharding_spec(self, dim_partition_list):\n\n        sharding_spec_list = []\n        check_duplicated_list = []\n        for output_dim_partition_dict in dim_partition_list:\n            try:\n                output_sharding_spec = self._generate_sharding_spec(self.output_data, output_dim_partition_dict)\n            except AssertionError as e:\n                warnings.warn(f'{e}')\n                break\n            sharding_seq = output_sharding_spec.sharding_sequence\n            if sharding_seq not in check_duplicated_list:\n                check_duplicated_list.append(sharding_seq)\n                sharding_spec_list.append(output_sharding_spec)\n\n        return sharding_spec_list\n\n    def _enumerate_all_possible_output(self, mesh_dim_0, mesh_dim_1):\n        # use mesh_dim_0, mesh_dim_1 instead of constant 0, 1 in here for N-D device mesh scaliablity.\n\n        output_dim_partition_list = []\n        dim_size = self.output_data.dim()\n        # enumerate all the 2D sharding cases\n        sharding_list_2d = enumerate_all_possible_2d_sharding(mesh_dim_0, mesh_dim_1, dim_size)\n        output_dim_partition_list.extend(sharding_list_2d)\n\n        # enumerate all the 1D sharding cases\n        sharding_list_1d_on_dim_0 = enumerate_all_possible_1d_sharding(mesh_dim_0, dim_size)\n        output_dim_partition_list.extend(sharding_list_1d_on_dim_0)\n        sharding_list_1d_on_dim_1 = enumerate_all_possible_1d_sharding(mesh_dim_1, dim_size)\n        output_dim_partition_list.extend(sharding_list_1d_on_dim_1)\n\n        # add empty dict for fully replicated case\n        output_dim_partition_list.append({})\n        output_sharding_spec_list = self._convert_partition_dict_to_sharding_spec(output_dim_partition_list)\n\n        return output_sharding_spec_list\n\n    @ignore_sharding_exception\n    def _register_strategy(self, output_sharding_spec):\n        dim_partition_dict_for_input = output_sharding_spec.dim_partition_dict\n        sharding_spec_for_condition = self._generate_sharding_spec(self.condition_data, dim_partition_dict_for_input)\n        sharding_spec_for_x = self._generate_sharding_spec(self.x_data, dim_partition_dict_for_input)\n        sharding_spec_for_y = self._generate_sharding_spec(self.y_data, dim_partition_dict_for_input)\n\n        name = f'{output_sharding_spec.sharding_sequence} = {sharding_spec_for_condition.sharding_sequence} x {sharding_spec_for_x.sharding_sequence} x {sharding_spec_for_y.sharding_sequence}'\n        dim_partition_dict_for_output = output_sharding_spec.dim_partition_dict\n\n        # generate resharding cost for this strategy\n        resharding_costs = self._generate_resharding_costs(\n            [sharding_spec_for_condition, sharding_spec_for_x, sharding_spec_for_y])\n\n        # compute the computation cost of this strategy\n        sharding_dims = []\n        for mesh_dims in dim_partition_dict_for_output.values():\n            for mesh_dim in mesh_dims:\n                sharding_dims.append(self.device_mesh.shape[mesh_dim])\n        sharding_size = reduce(operator.mul, sharding_dims, 1)\n        memory_cost = self.output_data.numel() / sharding_size\n        compute_cost = memory_cost\n        communication_cost = 0\n\n        sharding_strategies = ShardingStrategy(name,\n                                               output_sharding_spec=output_sharding_spec,\n                                               compute_cost=compute_cost,\n                                               communication_cost=communication_cost,\n                                               memory_cost=memory_cost,\n                                               resharding_costs=resharding_costs,\n                                               input_shardings=(sharding_spec_for_condition, sharding_spec_for_x,\n                                                                sharding_spec_for_y))\n\n        self.strategies_vector.append(sharding_strategies)\n\n    def register_strategy(self) -> StrategiesVector:\n        MESH_DIM_LIST = [0, 1]\n        output_sharding_specs = self._enumerate_all_possible_output(MESH_DIM_LIST[0], MESH_DIM_LIST[1])\n        for output_sharding_spec in output_sharding_specs:\n            self._register_strategy(output_sharding_spec)\n", "path": "colossalai/auto_parallel/tensor_shard/deprecated/op_handler/where_handler.py"}]}