code: def _broadcast_bmm(a, b) -> torch.Tensor:
      """
      Batch multiply two matrices and broadcast if necessary.
  
      Args:
          a: torch tensor of shape (P, K) or (M, P, K)
          b: torch tensor of shape (N, K, K)
  
      Returns:
          a and b broadcast multiplied. The output batch dimension is max(N, M).
  
      To broadcast transforms across a batch dimension if M != N then
      expect that either M = 1 or N = 1. The tensor with batch dimension 1 is
      expanded to have shape N or M.
      """
      if a.dim() == 2:
          a = a[None]
      if len(a) != len(b):
          if not ((len(a) == 1) or (len(b) == 1)):
              msg = "Expected batch dim for bmm to be equal or 1; got %r, %r"
              raise ValueError(msg % (a.shape, b.shape))
          if len(a) == 1:
              a = a.expand(len(b), -1, -1)
          if len(b) == 1:
              b = b.expand(len(a), -1, -1)
      return a.bmm(b)
  
imports: import math
  import os
  import warnings
  from typing import List, Optional, Union
  import torch
  from ..common.datatypes import Device, get_device, make_device
  from ..common.workaround import _safe_det_3x3
  from .rotation_conversions import _axis_angle_rotation
  from .se3 import se3_log_map
label_desc: Use image node_id 4 for calling a global function with example usage: _broadcast_bmm(a, b) and returns: a