code: def cat2(lhs, rhs, *, dim=-1):
      """
      Like ``torch.cat([lhs, rhs], dim=dim)`` but dispatches to
      :func:`torch.nn.functional.pad` in case one of ``lhs`` or ``rhs`` is a
      scalar.
      """
      assert dim < 0
      if not isinstance(lhs, torch.Tensor):
          pad = (0, 0) * (-1 - dim) + (1, 0)
          return torch.nn.functional.pad(rhs, pad, value=lhs)
      if not isinstance(rhs, torch.Tensor):
          pad = (0, 0) * (-1 - dim) + (0, 1)
          return torch.nn.functional.pad(lhs, pad, value=rhs)
  
      diff = lhs.dim() - rhs.dim()
      if diff > 0:
          rhs = rhs.expand((1,) * diff + rhs.shape)
      elif diff < 0:
          diff = -diff
          lhs = lhs.expand((1,) * diff + lhs.shape)
      shape = list(broadcast_shape(lhs.shape, rhs.shape))
      shape[dim] = -1
      return torch.cat([lhs.expand(shape), rhs.expand(shape)], dim=dim)
  
file_path: pyro/pyro/contrib/epidemiology/util.py
imports: import numpy
  import torch
  import pyro
  import pyro.distributions
  import pyro.poutine
  from pyro.distributions.util import broadcast_shape
  from pyro.ops.special import safe_log
label_desc: Use image node_id 2 for calling a global function with example usage: cat2(lhs, rhs) and returns: torch, torch, torch