code: def _get_data_mapping_for_bmm_op(
      node, input_idx, other_idx, bias_idx=None
  ):
      """
      This function is a helper function which extracts the common logic for both `bmm` and `addbmm`
      node handler to reduce code redundancy.
      """
      # input operand
      physical_input_operand = OperationData(
          name=str(node.args[input_idx]),
          type=OperationDataType.ARG,
          data=node.args[input_idx]._meta_data,
      )
  
      # other operand
      physical_other_operand = OperationData(
          name=str(node.args[other_idx]),
          type=OperationDataType.ARG,
          data=node.args[other_idx]._meta_data,
      )
  
      # output
      physical_output = OperationData(
          name=str(node),
          type=OperationDataType.OUTPUT,
          data=node._meta_data,
      )
      mapping = {
          "input": physical_input_operand,
          "other": physical_other_operand,
          "output": physical_output,
      }
  
      if bias_idx is not None:
          # bias physical shape
          bias_logical_shape = node._meta_data.shape
          physical_bias_operand = OperationData(
              name=str(node.args[bias_idx]),
              type=OperationDataType.ARG,
              data=node.args[bias_idx]._meta_data,
              logical_shape=bias_logical_shape,
          )
          mapping["bias"] = physical_bias_operand
      return mapping
  
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 1 for calling a global function with example usage: _get_data_mapping_for_bmm_op(node, input_idx, other_idx, bias_idx) and returns: mapping