prompt: 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()
  
  ```
  
golden_diff: 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
  
verification_info: {"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"}]}