Make attention weights non-0 only on the top-hard_attention_k ones.

def harden_attention_weights(weights, hard_attention_k):
  """Make attention weights non-0 only on the top-hard_attention_k ones."""
  # Subtract the top-kth weight and zero-out all lower ones.
  # Note that currently in case of numerical ties it will retain more
  # than k elements. In the future, we may want to avoid this.
  weights -= common_layers.top_kth_iterative(weights, hard_attention_k)
  weights = tf.nn.relu(weights)
  # Re-normalize the weights.
  weights_sum = tf.reduce_sum(weights, axis=-1, keep_dims=True)
  weights_sum = tf.maximum(weights_sum, 1e-6)  # Avoid division by 0.
  weights /= weights_sum
  return weights