code: def get_broadcast_shape(
      shape1: torch.Size, shape2: torch.Size
  ) -> List[int]:
      """
      Compute the broadcast shape given two shapes.
      """
      assert is_broadcastable(
          shape1, shape2
      ), f"{shape1} and {shape2} are not broadcastable"
      shape1_reverse = shape1[::-1]
      shape2_reverse = shape2[::-1]
      min_common_dim = min(len(shape1), len(shape2))
      dims = []
      for s1, s2 in zip(shape1_reverse, shape2_reverse):
          dims.append(max(s1, s2))
  
      # append the remaining dims
      dims.extend(shape1_reverse[min_common_dim:])
      dims.extend(shape2_reverse[min_common_dim:])
      return dims[::-1]
  
file_path: ColossalAI/colossalai/auto_parallel/tensor_shard/utils/broadcast.py
imports: from enum import Enum, auto
  from typing import List
  import torch
  from torch.fx.node import Node
  from colossalai.auto_parallel.tensor_shard.sharding_strategy import CommAction, CommType, OperationData, OperationDataType
  from colossalai.tensor.comm_spec import CollectiveCommPattern, CommSpec
  from colossalai.tensor.sharding_spec import ShardingSpec
label_desc: Use image node_id 2 for calling a global function with example usage: get_broadcast_shape(shape1, shape2) and returns: dims