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 recover(
      self,
      op_data_mapping: Dict[str, OperationData],
      strategy: ShardingStrategy,
  ):
      # remove sharding on the broadcast dim
      def _remove_sharding_on_broadcast_dim(key, strategy):
          op_data = op_data_mapping[key]
          sharding_spec = strategy.get_sharding_spec_by_name(
              op_data.name
          )
          tensor_shape = list(sharding_spec.entire_shape)
  
          for dim_idx, broadcast_type in self.broadcast_dim_info[
              key
          ].items():
              if broadcast_type == BroadcastType.MULTIPLE:
                  # if the dim is originally 1 and multiplied during broadcast
                  # we set its sharding to R
                  # e.g. [1, 2, 4] x [4, 4, 8] -> [4, 2, 8]
                  # the dim 0 of [1, 2, 4] is multiplied to 4
                  tensor_shape[dim_idx] = 1
              elif broadcast_type == BroadcastType.PADDING:
                  # if the dim is padded
                  # we remove its sharding
                  tensor_shape[dim_idx] = None
  
          tensor_shape_before_broadcast = [
              dim for dim in tensor_shape if dim is not None
          ]
  
          (
              physical_sharding_spec,
              removed_dims,
          ) = recover_sharding_spec_for_broadcast_shape(
              logical_sharding_spec=sharding_spec,
              logical_shape=sharding_spec.entire_shape,
              physical_shape=tensor_shape_before_broadcast,
          )
          strategy.sharding_specs[op_data] = physical_sharding_spec
  
      # enumerate all sharding strategies
      strategies = []
      try:
          strategy_copy = strategy.clone()
          _remove_sharding_on_broadcast_dim("input", strategy_copy)
          _remove_sharding_on_broadcast_dim("other", strategy_copy)
          strategies.append(strategy_copy)
      except ShardingSpecException:
          pass
      return strategies
  
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 3 for calling the Broadcaster obj's underlying member method code with example usage: obj.recover(op_data_mapping, strategy) and returns: strategies