class_docstr: Broadcast the non-matrix dimensions for batched matrix multiplication.
class_docstr_tok: ["Broadcast","the","non-matrix","dimensions","for","batched","matrix","multiplication","."]
code: def apply(self, shape_mapping: Dict[str, List[int]]):
      mapping_copy = shape_mapping.copy()
  
      # get shapes
      input_shape = mapping_copy["input"]
      other_shape = mapping_copy["other"]
  
      # sanity check
      assert len(input_shape) > 1 and len(other_shape) > 1
  
      # broadcast the batch dim and record
      bcast_non_matrix_dims = get_broadcast_shape(
          input_shape[:-2], other_shape[:-2]
      )
  
      # store the broadcast dim info
      input_broadcast_dim_info = get_broadcast_dim_info(
          bcast_non_matrix_dims, input_shape[:-2]
      )
      other_broadcast_dim_info = get_broadcast_dim_info(
          bcast_non_matrix_dims, other_shape[:-2]
      )
      self.broadcast_dim_info["input"] = input_broadcast_dim_info
      self.broadcast_dim_info["other"] = other_broadcast_dim_info
  
      # create the full logical shape
      input_shape = bcast_non_matrix_dims + input_shape[-2:]
      other_shape = bcast_non_matrix_dims + other_shape[-2:]
      assert len(input_shape) == len(other_shape)
  
      mapping_copy["input"] = input_shape
      mapping_copy["other"] = other_shape
  
      return mapping_copy
  
imports: import operator
  from abc import ABC, abstractmethod
  from copy import deepcopy
  from enum import Enum
  from functools import reduce
  from typing import Dict, List, Union
  import torch
  from colossalai.auto_parallel.tensor_shard.utils.broadcast import BroadcastType, get_broadcast_dim_info, get_broadcast_shape
  from colossalai.tensor.sharding_spec import ShardingSpecException
  from ..sharding_strategy import OperationData, OperationDataType, ShardingStrategy
  from ..utils import recover_sharding_spec_for_broadcast_shape
  from .node_handler import MetaInfoNodeHandler
  from .registry import operator_registry
  from .strategy import BatchedMatMulStrategyGenerator, DotProductStrategyGenerator, LinearProjectionStrategyGenerator, MatVecStrategyGenerator, StrategyGenerator
label_desc: Use image node_id 2 for calling the Broadcaster obj's underlying member method code with example usage: obj.apply(shape_mapping) and returns: mapping_copy