class_docstr: This is a NodeHandler class which deals with the addition + batched matrix multiplication operation in PyTorch.
  Such operations including `torch.addbmm` and `torch.Tensor.addbmm` require the two matmul tensor to be 3D. However, due to the
  addition, logical-physical shape conversion is required for the bias term.
  
  As the addbmm operation will reduce the batch dimension, the bias is maximum 2D.
class_docstr_tok: ["This","is","a","NodeHandler","class","which","deals","with","the","addition","+","batched","matrix","multiplication","operation","in","PyTorch",".","Such","operations","including","`","torch.addbmm","`","and","`","torch.Tensor.addbmm","`","require","the","two","matmul","tensor","to","be","3D",".","However",",","due","to","the","addition",",","logical-physical","shape","conversion","is","required","for","the","bias","term",".","As","the","addbmm","operation","will","reduce","the","batch","dimension",",","the","bias","is","maximum","2D","."]
code: def post_process(
      self, strategy: ShardingStrategy
  ) -> Union[ShardingStrategy, List[ShardingStrategy]]:
      # convert bias from its logical sharding spec to its physical sharding spec
      op_data_mapping = self.get_operation_data_mapping()
  
      if "bias" in op_data_mapping:
          bias_op_data = op_data_mapping["bias"]
          bias_physical_shape = bias_op_data.data.shape
          bias_logical_shape = bias_op_data.logical_shape
          bias_sharding_spec = strategy.get_sharding_spec_by_name(
              bias_op_data.name
          )
          (
              bias_sharding_spec,
              removed_dims,
          ) = recover_sharding_spec_for_broadcast_shape(
              bias_sharding_spec,
              bias_logical_shape,
              bias_physical_shape,
          )
          strategy.sharding_specs[bias_op_data] = bias_sharding_spec
  
          if len(removed_dims) > 0:
              comm_action = comm_actions_for_oprands(
                  node=self.node,
                  removed_dims=removed_dims,
                  op_data=bias_op_data,
                  sharding_spec=bias_sharding_spec,
              )
              strategy.communication_actions[bias_op_data] = comm_action
  
      return strategy
  
file_path: colossalai/colossalai/auto_parallel/tensor_shard/node_handler/bmm_handler.py
imports: from typing import Dict, List, Union
  import torch
  from ..sharding_strategy import OperationData, OperationDataType, ShardingStrategy
  from ..utils import comm_actions_for_oprands, recover_sharding_spec_for_broadcast_shape
  from .node_handler import NodeHandler
  from .registry import operator_registry
  from .strategy import BatchedMatMulStrategyGenerator, StrategyGenerator
label_desc: Use image node_id 3 for calling the AddBMMFunctionHandler obj's underlying member method code with example usage: obj.post_process(strategy) and returns: strategy