def _compute_attention(attention_mechanism, cell_output, attention_state,
                       attention_layer):
  """Computes the attention and alignments for a given attention_mechanism."""
  alignments, next_attention_state = attention_mechanism(
      cell_output, state=attention_state)

  # Reshape from [1, memory_time, memory_size] to [memory_time, memory_size]
  memory = array_ops.squeeze(attention_mechanism.values, 0)
  # Context is the inner product of alignments and values along the
  # memory time dimension.
  # alignments shape is
  #   [batch_size, memory_time]
  # memory shape is
  #   [memory_time, memory_size]
  # the matmul is over memory_time, so the output shape is
  #   [batch_size, memory_size].
  context = math_ops.matmul(alignments, memory)

  if attention_layer is not None:
    attention = attention_layer(array_ops.concat([cell_output, context], 1))
  else:
    attention = context

  return attention, alignments, next_attention_state
Computes the attention and alignments for a given attention_mechanism.
language/labs/consistent_zero_shot_nmt/modules/attention_wrappers.py