def _combine_sup_unsup_datasets(sup_data, unsup_data):
  """Combines supervised and usupervised samples into single dictionary.

  Args:
    sup_data: dictionary with examples from supervised dataset.
    unsup_data: dictionary with examples from unsupervised dataset.

  Returns:
    Dictionary with combined suvervised and unsupervised examples.
  """
  # Copy all values from supervised data as is
  output_dict = dict(sup_data)

  # take only 'image' and 'aug_image' from unsupervised dataset and
  # rename then into 'unsup_image' and 'unsup_aug_image'
  if 'image' in unsup_data:
    output_dict['unsup_image'] = unsup_data.pop('image')
  if 'aug_image' in unsup_data:
    output_dict['unsup_aug_image'] = unsup_data.pop('aug_image')

  return output_dict