class_docstr: This is a NodeHandler class which deals with the batched matrix multiplication operation in PyTorch.
  Such operations including `torch.bmm` and `torch.Tensor.bmm` require the tensor to be 3D, thus, there is
  no logical-physical shape conversion in this handler.
class_docstr_tok: ["This","is","a","NodeHandler","class","which","deals","with","the","batched","matrix","multiplication","operation","in","PyTorch",".","Such","operations","including","`","torch.bmm","`","and","`","torch.Tensor.bmm","`","require","the","tensor","to","be","3D",",","thus",",","there","is","no","logical-physical","shape","conversion","in","this","handler","."]
code: def get_operation_data_mapping(self) -> Dict[str, OperationData]:
      # input operand
      input_data = self.node.args[1]._meta_data
      physical_input_operand = OperationData(
          name=str(self.node.args[1]),
          type=self._infer_op_data_type(input_data),
          data=input_data,
      )
  
      # other operand
      other_data = self.node.args[2]._meta_data
      physical_other_operand = OperationData(
          name=str(self.node.args[2]),
          type=self._infer_op_data_type(other_data),
          data=other_data,
      )
      # bias physical shape
      bias_logical_shape = self.node._meta_data.shape
      bias_data = self.node.args[0]._meta_data
      physical_bias_operand = OperationData(
          name=str(self.node.args[0]),
          type=self._infer_op_data_type(bias_data),
          data=bias_data,
          logical_shape=bias_logical_shape,
      )
  
      # output
      physical_output = OperationData(
          name=str(self.node),
          type=OperationDataType.OUTPUT,
          data=self.node._meta_data,
      )
  
      mapping = {
          "input": physical_input_operand,
          "other": physical_other_operand,
          "output": physical_output,
          "bias": physical_bias_operand,
      }
  
      return mapping
  
file_path: ColossalAI/colossalai/auto_parallel/tensor_shard/node_handler/addmm_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 LinearProjectionStrategyGenerator, StrategyGenerator
label_desc: Use image node_id 2 for calling the ADDMMFunctionHandler obj's underlying member method code with example usage: obj.get_operation_data_mapping() and returns: mapping