def stack_embeddings(model_outputs, embedding_keys, common_module=common):
  """Selects and stacks embeddings by key.

  Args:
    model_outputs: A dictionary for model output tensors.
    embedding_keys: A list for enum strings for tensor keys to select.
    common_module: A Python module that defines common flags and constants.

  Returns:
    A tensor for stacked embeddings. Shape = [..., num_embeddings_per_instance,
      embedding_dim].
  """
  embeddings_to_stack = []
  for key in embedding_keys:
    if key in [
        common_module.KEY_EMBEDDING_MEANS, common_module.KEY_EMBEDDING_STDDEVS
    ]:
      embeddings_to_stack.append(tf.expand_dims(model_outputs[key], axis=-2))
    elif key == common_module.KEY_EMBEDDING_SAMPLES:
      embeddings_to_stack.append(model_outputs[key])
    else:
      raise ValueError('Unsupported embedding key: `%s`.' % str(key))
  return tf.concat(embeddings_to_stack, axis=-2)