import torch


def rgba_premultiply(rgba):
  """ Premultiplies the RGB channels with the Alpha channel
    Args:
        rgba: [batch, channels=4, height, width] range [0, 1]
    Returns:
        returns [batch,channels=4, height, width] range [0, 1]
  """
  premultiplied_rgb = rgba[:,:3,:,:] * rgba[:,3,:,:].unsqueeze(1).expand(-1, 3, -1, -1)
  return torch.cat([premultiplied_rgb, rgba[:,3,:,:].unsqueeze(1)], 1)