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

[FEATURE]: Patch meta information of `torch.nn.functional.softmax()`
This is a part of issue #2628, we will patch meta information of `torch.nn.functional.softmax()`

--- 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/meta_profiler/meta_registry/activation.py`
Content:
```
1 from typing import List, Tuple
2 
3 import torch
4 
5 from colossalai.auto_parallel.tensor_shard.sharding_strategy import MemoryCost, OperationDataType, TrainCycleItem
6 from colossalai.fx.profiler.memory_utils import activation_size
7 from colossalai.fx.profiler.opcount import flop_mapping
8 
9 from ..registry import meta_register
10 
11 __all__ = ["relu_meta_info"]
12 
13 
14 @meta_register.register(torch.nn.ReLU)
15 def relu_meta_info(*args, **kwargs) -> Tuple[TrainCycleItem, TrainCycleItem, List[torch.Tensor]]:
16     """torch.nn.ReLU metainfo generator
17     The aten graph of torch.nn.ReLU is
18     graph():
19     %input_2 : [#users=1] = placeholder[target=placeholder](default=)
20     %relu_default : [#users=2] = call_function[target=torch.ops.aten.relu.default](args = (%input_2,), kwargs = {})
21     %zeros_like_default : [#users=1] = call_function[target=torch.ops.aten.zeros_like.default](args = (%relu_default,), kwargs = {dtype: None, layout: None, device: None, pin_memory: None})
22     %detach_default : [#users=1] = call_function[target=torch.ops.aten.detach.default](args = (%relu_default,), kwargs = {})
23     %threshold_backward_default : [#users=1] = call_function[target=torch.ops.aten.threshold_backward.default](args = (%zeros_like_default, %detach_default, None), kwargs = {})
24     %detach_default_1 : [#users=1] = call_function[target=torch.ops.aten.detach.default](args = (%threshold_backward_default,), kwargs = {})
25     %detach_default_2 : [#users=0] = call_function[target=torch.ops.aten.detach.default](args = (%detach_default_1,), kwargs = {})
26 
27     Returns:
28         Tuple[TrainCycleItem, TrainCycleItem, List[torch.Tensor]]: compute cost, memory cost and forward inputs
29     """
30 
31     input_tensor = args[0].data
32     output_tensor = next(filter(lambda x: x.type == OperationDataType.OUTPUT, args)).data
33     is_inplace = kwargs.get("inplace", False)
34 
35     # construct input args for forward
36     fwd_in_args = [input_tensor]
37 
38     # construct input args for backward
39     bwd_in_args = [output_tensor]
40 
41     # calculate cost
42     # the fwd op with compute cost is relu.default
43     # the bwd op with compute cost is threshold_backward
44 
45     # calculate compute cost
46     fwd_compute_cost = flop_mapping[torch.ops.aten.relu.default](fwd_in_args, (output_tensor,))
47     bwd_compute_cost = flop_mapping[torch.ops.aten.threshold_backward.default](bwd_in_args, (input_tensor,))
48     compute_cost = TrainCycleItem(fwd=fwd_compute_cost, bwd=bwd_compute_cost, total=fwd_compute_cost + bwd_compute_cost)
49 
50     # calculate memory cost
51     # NOTE: the inplace ReLU don't have forward memory cost
52     # NOTE: currently in SPMD solver we always believe that there will be a new tensor created in forward
53     fwd_memory_cost = MemoryCost(
54         activation=activation_size(input_tensor) if is_inplace else activation_size([output_tensor, input_tensor]),
55         parameter=0,
56         temp=0,
57         buffer=0)
58 
59     bwd_memory_cost = MemoryCost(activation=activation_size(input_tensor), parameter=0, temp=0, buffer=0)
60 
61     # total cost is the sum of forward and backward cost
62     total_cost = MemoryCost(activation=fwd_memory_cost.activation + bwd_memory_cost.activation,
63                             parameter=fwd_memory_cost.parameter + bwd_memory_cost.parameter)
64 
65     memory_cost = TrainCycleItem(fwd=fwd_memory_cost, bwd=bwd_memory_cost, total=total_cost)
66 
67     # store fwd_in, fwd_buffer, fwd_out
68     # NOTE: It might seems a little bit weird here, we just want to align it with the older version
69     # of MetaInfoProp. In the future we might modify this part to make it clearer.
70     fwd_in = []
71     fwd_buffer = [torch.zeros_like(output_tensor, device='meta')]
72     fwd_out = [torch.zeros_like(output_tensor, device='meta')]
73 
74     return compute_cost, memory_cost, fwd_in, fwd_buffer, fwd_out
75
```
--- 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/meta_profiler/meta_registry/activation.py b/colossalai/auto_parallel/meta_profiler/meta_registry/activation.py
--- a/colossalai/auto_parallel/meta_profiler/meta_registry/activation.py
+++ b/colossalai/auto_parallel/meta_profiler/meta_registry/activation.py
@@ -72,3 +72,53 @@
     fwd_out = [torch.zeros_like(output_tensor, device='meta')]
 
     return compute_cost, memory_cost, fwd_in, fwd_buffer, fwd_out
+
+
+@meta_register.register(torch.nn.Softmax)
+@meta_register.register(torch.nn.functional.softmax)
+def softmax_meta_info(*args, **kwargs) -> Tuple[TrainCycleItem, TrainCycleItem, List[torch.Tensor]]:
+    """torch.nn.Softmax metainfo generator
+    Returns:
+        Tuple[TrainCycleItem, TrainCycleItem, List[torch.Tensor]]: compute cost, memory cost and forward inputs
+    """
+    input_tensor = next(
+        filter(
+            lambda x:
+            (x.type == OperationDataType.ARG or x.type == OperationDataType.PARAM) and x.name != 'softmax_dim',
+            args)).data
+    output_tensor = next(filter(lambda x: x.type == OperationDataType.OUTPUT, args)).data
+    softmax_dim = next(filter(lambda x: x.name == 'softmax_dim', args)).data
+
+    # calculate cost
+
+    # calculate compute cost
+    fwd_compute_cost = flop_mapping[torch.ops.aten._softmax.default]([input_tensor], [output_tensor])
+    bwd_compute_cost = flop_mapping[torch.ops.aten._softmax_backward_data.default]([output_tensor], [input_tensor])
+
+    compute_cost = TrainCycleItem(fwd=fwd_compute_cost, bwd=bwd_compute_cost, total=fwd_compute_cost + bwd_compute_cost)
+
+    # calculate memory cost
+    # NOTE: currently in SPMD solver we always believe that there will be a new tensor created in forward
+    fwd_memory_cost = MemoryCost(activation=activation_size([input_tensor, output_tensor]),
+                                 parameter=0,
+                                 temp=0,
+                                 buffer=0)
+    bwd_memory_cost = MemoryCost(activation=activation_size(input_tensor),
+                                 parameter=0,
+                                 temp=activation_size(input_tensor),
+                                 buffer=0)
+
+    # total cost is the sum of forward and backward cost
+    total_cost = MemoryCost(activation=fwd_memory_cost.activation + bwd_memory_cost.activation,
+                            parameter=fwd_memory_cost.parameter + bwd_memory_cost.parameter,
+                            temp=fwd_memory_cost.temp + bwd_memory_cost.temp,
+                            buffer=fwd_memory_cost.buffer + bwd_memory_cost.buffer)
+
+    memory_cost = TrainCycleItem(fwd=fwd_memory_cost, bwd=bwd_memory_cost, total=total_cost)
+
+    # store fwd_in, fwd_buffer, fwd_out
+    fwd_in = []
+    fwd_buffer = [torch.zeros_like(output_tensor, device='meta')]
+    fwd_out = [torch.zeros_like(output_tensor, device='meta')]
+
+    return compute_cost, memory_cost, fwd_in, fwd_buffer, fwd_out

{"golden_diff": "diff --git a/colossalai/auto_parallel/meta_profiler/meta_registry/activation.py b/colossalai/auto_parallel/meta_profiler/meta_registry/activation.py\n--- a/colossalai/auto_parallel/meta_profiler/meta_registry/activation.py\n+++ b/colossalai/auto_parallel/meta_profiler/meta_registry/activation.py\n@@ -72,3 +72,53 @@\n     fwd_out = [torch.zeros_like(output_tensor, device='meta')]\n \n     return compute_cost, memory_cost, fwd_in, fwd_buffer, fwd_out\n+\n+\n+@meta_register.register(torch.nn.Softmax)\n+@meta_register.register(torch.nn.functional.softmax)\n+def softmax_meta_info(*args, **kwargs) -> Tuple[TrainCycleItem, TrainCycleItem, List[torch.Tensor]]:\n+    \"\"\"torch.nn.Softmax metainfo generator\n+    Returns:\n+        Tuple[TrainCycleItem, TrainCycleItem, List[torch.Tensor]]: compute cost, memory cost and forward inputs\n+    \"\"\"\n+    input_tensor = next(\n+        filter(\n+            lambda x:\n+            (x.type == OperationDataType.ARG or x.type == OperationDataType.PARAM) and x.name != 'softmax_dim',\n+            args)).data\n+    output_tensor = next(filter(lambda x: x.type == OperationDataType.OUTPUT, args)).data\n+    softmax_dim = next(filter(lambda x: x.name == 'softmax_dim', args)).data\n+\n+    # calculate cost\n+\n+    # calculate compute cost\n+    fwd_compute_cost = flop_mapping[torch.ops.aten._softmax.default]([input_tensor], [output_tensor])\n+    bwd_compute_cost = flop_mapping[torch.ops.aten._softmax_backward_data.default]([output_tensor], [input_tensor])\n+\n+    compute_cost = TrainCycleItem(fwd=fwd_compute_cost, bwd=bwd_compute_cost, total=fwd_compute_cost + bwd_compute_cost)\n+\n+    # calculate memory cost\n+    # NOTE: currently in SPMD solver we always believe that there will be a new tensor created in forward\n+    fwd_memory_cost = MemoryCost(activation=activation_size([input_tensor, output_tensor]),\n+                                 parameter=0,\n+                                 temp=0,\n+                                 buffer=0)\n+    bwd_memory_cost = MemoryCost(activation=activation_size(input_tensor),\n+                                 parameter=0,\n+                                 temp=activation_size(input_tensor),\n+                                 buffer=0)\n+\n+    # total cost is the sum of forward and backward cost\n+    total_cost = MemoryCost(activation=fwd_memory_cost.activation + bwd_memory_cost.activation,\n+                            parameter=fwd_memory_cost.parameter + bwd_memory_cost.parameter,\n+                            temp=fwd_memory_cost.temp + bwd_memory_cost.temp,\n+                            buffer=fwd_memory_cost.buffer + bwd_memory_cost.buffer)\n+\n+    memory_cost = TrainCycleItem(fwd=fwd_memory_cost, bwd=bwd_memory_cost, total=total_cost)\n+\n+    # store fwd_in, fwd_buffer, fwd_out\n+    fwd_in = []\n+    fwd_buffer = [torch.zeros_like(output_tensor, device='meta')]\n+    fwd_out = [torch.zeros_like(output_tensor, device='meta')]\n+\n+    return compute_cost, memory_cost, fwd_in, fwd_buffer, fwd_out\n", "issue": "[tensor] fix some unittests\n\n[tensor] fix some unittests\n\n[tensor] fix some unittests\n\n[FEATURE]: Patch meta information of `torch.nn.functional.softmax()`\nThis is a part of issue #2628, we will patch meta information of `torch.nn.functional.softmax()`\n", "before_files": [{"content": "from typing import List, Tuple\n\nimport torch\n\nfrom colossalai.auto_parallel.tensor_shard.sharding_strategy import MemoryCost, OperationDataType, TrainCycleItem\nfrom colossalai.fx.profiler.memory_utils import activation_size\nfrom colossalai.fx.profiler.opcount import flop_mapping\n\nfrom ..registry import meta_register\n\n__all__ = [\"relu_meta_info\"]\n\n\n@meta_register.register(torch.nn.ReLU)\ndef relu_meta_info(*args, **kwargs) -> Tuple[TrainCycleItem, TrainCycleItem, List[torch.Tensor]]:\n    \"\"\"torch.nn.ReLU metainfo generator\n    The aten graph of torch.nn.ReLU is\n    graph():\n    %input_2 : [#users=1] = placeholder[target=placeholder](default=)\n    %relu_default : [#users=2] = call_function[target=torch.ops.aten.relu.default](args = (%input_2,), kwargs = {})\n    %zeros_like_default : [#users=1] = call_function[target=torch.ops.aten.zeros_like.default](args = (%relu_default,), kwargs = {dtype: None, layout: None, device: None, pin_memory: None})\n    %detach_default : [#users=1] = call_function[target=torch.ops.aten.detach.default](args = (%relu_default,), kwargs = {})\n    %threshold_backward_default : [#users=1] = call_function[target=torch.ops.aten.threshold_backward.default](args = (%zeros_like_default, %detach_default, None), kwargs = {})\n    %detach_default_1 : [#users=1] = call_function[target=torch.ops.aten.detach.default](args = (%threshold_backward_default,), kwargs = {})\n    %detach_default_2 : [#users=0] = call_function[target=torch.ops.aten.detach.default](args = (%detach_default_1,), kwargs = {})\n\n    Returns:\n        Tuple[TrainCycleItem, TrainCycleItem, List[torch.Tensor]]: compute cost, memory cost and forward inputs\n    \"\"\"\n\n    input_tensor = args[0].data\n    output_tensor = next(filter(lambda x: x.type == OperationDataType.OUTPUT, args)).data\n    is_inplace = kwargs.get(\"inplace\", False)\n\n    # construct input args for forward\n    fwd_in_args = [input_tensor]\n\n    # construct input args for backward\n    bwd_in_args = [output_tensor]\n\n    # calculate cost\n    # the fwd op with compute cost is relu.default\n    # the bwd op with compute cost is threshold_backward\n\n    # calculate compute cost\n    fwd_compute_cost = flop_mapping[torch.ops.aten.relu.default](fwd_in_args, (output_tensor,))\n    bwd_compute_cost = flop_mapping[torch.ops.aten.threshold_backward.default](bwd_in_args, (input_tensor,))\n    compute_cost = TrainCycleItem(fwd=fwd_compute_cost, bwd=bwd_compute_cost, total=fwd_compute_cost + bwd_compute_cost)\n\n    # calculate memory cost\n    # NOTE: the inplace ReLU don't have forward memory cost\n    # NOTE: currently in SPMD solver we always believe that there will be a new tensor created in forward\n    fwd_memory_cost = MemoryCost(\n        activation=activation_size(input_tensor) if is_inplace else activation_size([output_tensor, input_tensor]),\n        parameter=0,\n        temp=0,\n        buffer=0)\n\n    bwd_memory_cost = MemoryCost(activation=activation_size(input_tensor), parameter=0, temp=0, buffer=0)\n\n    # total cost is the sum of forward and backward cost\n    total_cost = MemoryCost(activation=fwd_memory_cost.activation + bwd_memory_cost.activation,\n                            parameter=fwd_memory_cost.parameter + bwd_memory_cost.parameter)\n\n    memory_cost = TrainCycleItem(fwd=fwd_memory_cost, bwd=bwd_memory_cost, total=total_cost)\n\n    # store fwd_in, fwd_buffer, fwd_out\n    # NOTE: It might seems a little bit weird here, we just want to align it with the older version\n    # of MetaInfoProp. In the future we might modify this part to make it clearer.\n    fwd_in = []\n    fwd_buffer = [torch.zeros_like(output_tensor, device='meta')]\n    fwd_out = [torch.zeros_like(output_tensor, device='meta')]\n\n    return compute_cost, memory_cost, fwd_in, fwd_buffer, fwd_out\n", "path": "colossalai/auto_parallel/meta_profiler/meta_registry/activation.py"}], "after_files": [{"content": "from typing import List, Tuple\n\nimport torch\n\nfrom colossalai.auto_parallel.tensor_shard.sharding_strategy import MemoryCost, OperationDataType, TrainCycleItem\nfrom colossalai.fx.profiler.memory_utils import activation_size\nfrom colossalai.fx.profiler.opcount import flop_mapping\n\nfrom ..registry import meta_register\n\n__all__ = [\"relu_meta_info\"]\n\n\n@meta_register.register(torch.nn.ReLU)\ndef relu_meta_info(*args, **kwargs) -> Tuple[TrainCycleItem, TrainCycleItem, List[torch.Tensor]]:\n    \"\"\"torch.nn.ReLU metainfo generator\n    The aten graph of torch.nn.ReLU is\n    graph():\n    %input_2 : [#users=1] = placeholder[target=placeholder](default=)\n    %relu_default : [#users=2] = call_function[target=torch.ops.aten.relu.default](args = (%input_2,), kwargs = {})\n    %zeros_like_default : [#users=1] = call_function[target=torch.ops.aten.zeros_like.default](args = (%relu_default,), kwargs = {dtype: None, layout: None, device: None, pin_memory: None})\n    %detach_default : [#users=1] = call_function[target=torch.ops.aten.detach.default](args = (%relu_default,), kwargs = {})\n    %threshold_backward_default : [#users=1] = call_function[target=torch.ops.aten.threshold_backward.default](args = (%zeros_like_default, %detach_default, None), kwargs = {})\n    %detach_default_1 : [#users=1] = call_function[target=torch.ops.aten.detach.default](args = (%threshold_backward_default,), kwargs = {})\n    %detach_default_2 : [#users=0] = call_function[target=torch.ops.aten.detach.default](args = (%detach_default_1,), kwargs = {})\n\n    Returns:\n        Tuple[TrainCycleItem, TrainCycleItem, List[torch.Tensor]]: compute cost, memory cost and forward inputs\n    \"\"\"\n\n    input_tensor = args[0].data\n    output_tensor = next(filter(lambda x: x.type == OperationDataType.OUTPUT, args)).data\n    is_inplace = kwargs.get(\"inplace\", False)\n\n    # construct input args for forward\n    fwd_in_args = [input_tensor]\n\n    # construct input args for backward\n    bwd_in_args = [output_tensor]\n\n    # calculate cost\n    # the fwd op with compute cost is relu.default\n    # the bwd op with compute cost is threshold_backward\n\n    # calculate compute cost\n    fwd_compute_cost = flop_mapping[torch.ops.aten.relu.default](fwd_in_args, (output_tensor,))\n    bwd_compute_cost = flop_mapping[torch.ops.aten.threshold_backward.default](bwd_in_args, (input_tensor,))\n    compute_cost = TrainCycleItem(fwd=fwd_compute_cost, bwd=bwd_compute_cost, total=fwd_compute_cost + bwd_compute_cost)\n\n    # calculate memory cost\n    # NOTE: the inplace ReLU don't have forward memory cost\n    # NOTE: currently in SPMD solver we always believe that there will be a new tensor created in forward\n    fwd_memory_cost = MemoryCost(\n        activation=activation_size(input_tensor) if is_inplace else activation_size([output_tensor, input_tensor]),\n        parameter=0,\n        temp=0,\n        buffer=0)\n\n    bwd_memory_cost = MemoryCost(activation=activation_size(input_tensor), parameter=0, temp=0, buffer=0)\n\n    # total cost is the sum of forward and backward cost\n    total_cost = MemoryCost(activation=fwd_memory_cost.activation + bwd_memory_cost.activation,\n                            parameter=fwd_memory_cost.parameter + bwd_memory_cost.parameter)\n\n    memory_cost = TrainCycleItem(fwd=fwd_memory_cost, bwd=bwd_memory_cost, total=total_cost)\n\n    # store fwd_in, fwd_buffer, fwd_out\n    # NOTE: It might seems a little bit weird here, we just want to align it with the older version\n    # of MetaInfoProp. In the future we might modify this part to make it clearer.\n    fwd_in = []\n    fwd_buffer = [torch.zeros_like(output_tensor, device='meta')]\n    fwd_out = [torch.zeros_like(output_tensor, device='meta')]\n\n    return compute_cost, memory_cost, fwd_in, fwd_buffer, fwd_out\n\n\n@meta_register.register(torch.nn.Softmax)\n@meta_register.register(torch.nn.functional.softmax)\ndef softmax_meta_info(*args, **kwargs) -> Tuple[TrainCycleItem, TrainCycleItem, List[torch.Tensor]]:\n    \"\"\"torch.nn.Softmax metainfo generator\n    Returns:\n        Tuple[TrainCycleItem, TrainCycleItem, List[torch.Tensor]]: compute cost, memory cost and forward inputs\n    \"\"\"\n    input_tensor = next(\n        filter(\n            lambda x:\n            (x.type == OperationDataType.ARG or x.type == OperationDataType.PARAM) and x.name != 'softmax_dim',\n            args)).data\n    output_tensor = next(filter(lambda x: x.type == OperationDataType.OUTPUT, args)).data\n    softmax_dim = next(filter(lambda x: x.name == 'softmax_dim', args)).data\n\n    # calculate cost\n\n    # calculate compute cost\n    fwd_compute_cost = flop_mapping[torch.ops.aten._softmax.default]([input_tensor], [output_tensor])\n    bwd_compute_cost = flop_mapping[torch.ops.aten._softmax_backward_data.default]([output_tensor], [input_tensor])\n\n    compute_cost = TrainCycleItem(fwd=fwd_compute_cost, bwd=bwd_compute_cost, total=fwd_compute_cost + bwd_compute_cost)\n\n    # calculate memory cost\n    # NOTE: currently in SPMD solver we always believe that there will be a new tensor created in forward\n    fwd_memory_cost = MemoryCost(activation=activation_size([input_tensor, output_tensor]),\n                                 parameter=0,\n                                 temp=0,\n                                 buffer=0)\n    bwd_memory_cost = MemoryCost(activation=activation_size(input_tensor),\n                                 parameter=0,\n                                 temp=activation_size(input_tensor),\n                                 buffer=0)\n\n    # total cost is the sum of forward and backward cost\n    total_cost = MemoryCost(activation=fwd_memory_cost.activation + bwd_memory_cost.activation,\n                            parameter=fwd_memory_cost.parameter + bwd_memory_cost.parameter,\n                            temp=fwd_memory_cost.temp + bwd_memory_cost.temp,\n                            buffer=fwd_memory_cost.buffer + bwd_memory_cost.buffer)\n\n    memory_cost = TrainCycleItem(fwd=fwd_memory_cost, bwd=bwd_memory_cost, total=total_cost)\n\n    # store fwd_in, fwd_buffer, fwd_out\n    fwd_in = []\n    fwd_buffer = [torch.zeros_like(output_tensor, device='meta')]\n    fwd_out = [torch.zeros_like(output_tensor, device='meta')]\n\n    return compute_cost, memory_cost, fwd_in, fwd_buffer, fwd_out\n", "path": "colossalai/auto_parallel/meta_profiler/meta_registry/activation.py"}]}