""" Implementations of attention layers.
"""

import tensorflow as tf
from seq2seq import GraphModule
from tensorflow.python.framework import function

class AttentionLayer(GraphModule):
  """
  Attention layer according to https://arxiv.org/abs/1409.0473.

  Args:
    num_units: Number of units used in the attention layer
    name: Name for this graph module
  """

  def __init__(self, num_units, name="attention"):
    super(AttentionLayer, self).__init__(name=name)
    self.num_units = num_units

  def _build(self, state, inputs):
    """Computes attention scores and outputs.

    Args:
      state: The state based on which to calculate attention scores.
        In seq2seq this is typically the current state of the decoder.
        A tensor of shape `[B, ...]`
      inputs: The elements to compute attention *over*. In seq2seq this is
        typically the sequence of encoder outputs.
        A tensor of shape `[B, T, input_dim]`

    Returns:
      A tuple `(scores, context)`.
      `scores` is vector of length `T` where each element is the
      normalized "score" of the corresponding `inputs` element.
      `context` is the final attention layer output corresponding to
      the weighted inputs.
      A tensor fo shape `[B, input_dim]`.
    """
    batch_size, inputs_timesteps, _ = tf.unpack(tf.shape(inputs))
    inputs_dim = inputs.get_shape().as_list()[-1]

    # Fully connected layers to transform both inputs and state
    # into a tensor with `num_units` units
    inputs_att = tf.contrib.layers.fully_connected(
        inputs=inputs,
        num_outputs=self.num_units,
        activation_fn=None,
        scope="inputs_att")
    state_att = tf.contrib.layers.fully_connected(
        inputs=state,
        num_outputs=self.num_units,
        activation_fn=None,
        scope="state_att")

    def bahdanau_score(inputs, state):
      """Computes Bahdanau-style attention scores
      """
      inputs.set_shape([None, None, self.num_units])
      state.set_shape([None, self.num_units])
      v_att = tf.get_variable("v_att", shape=[self.num_units], dtype=tf.float32)
      att_combined = tf.nn.tanh(inputs + tf.expand_dims(state, 1))
      att_combined_flat = tf.reshape(att_combined, [-1, self.num_units])
      scores = tf.matmul(att_combined_flat, tf.expand_dims(v_att, 1))
      scores = tf.reshape(scores, [batch_size, inputs_timesteps], name="scores")
      return scores

    scores = bahdanau_score(inputs_att, state_att)

    # Show, Attend, Spell type of attention
    # Take the dot product of state for each time step in inputs
    # Result: A tensor of shape [B, T]
    # inputs_att_flat = tf.reshape(inputs_att, [-1, self.num_units])
    # state_att_flat = tf.reshape(
    #     tf.tile(state_att, [1, inputs_timesteps]),
    #     [inputs_timesteps * batch_size, self.num_units])
    # scores = tf.matmul(
    #     tf.expand_dims(inputs_att_flat, 1), tf.expand_dims(state_att_flat, 2))
    # scores = tf.reshape(scores, [batch_size, inputs_timesteps], name="scores")

    # Normalize the scores
    scores_normalized = tf.nn.softmax(scores, name="scores_normalized")

    # Calculate the weighted average of the attention inputs
    # according to the scores
    context = tf.expand_dims(scores_normalized, 2) * inputs
    context = tf.reduce_sum(context, 1, name="context")
    context.set_shape([None, inputs_dim])

    return (scores_normalized, context)

@@ -43,32 +43,8 @@
 %22%22%0A%0A
-import tensorflow as tf%0A
 from
@@ -71,16 +71,40 @@
 hModule%0A
+import tensorflow as tf%0A
 from ten
@@ -1265,16 +1265,18 @@
  %22%22%22%0A   
+ #
  batch_s
@@ -1923,107 +1923,8 @@
 %22%22%22%0A
-      inputs.set_shape(%5BNone, None, self.num_units%5D)%0A      state.set_shape(%5BNone, self.num_units%5D)%0A
     
@@ -2004,318 +2004,386 @@
 32)%0A
+%0A
       
-att_combined = tf.nn.tanh(inputs + tf.expand_dims(state, 1))%0A      att_combined_flat = tf.reshape(att_combined, %5B-1, self.num_units%5D)%0A      scores = tf.matmul(att_combined_flat, tf.expand_dims(v_att, 1))%0A      scores = tf.reshape(scores, %5Bbatch_size, inputs_timesteps%5D, name=%22scores%22)%0A      return scores
+@function.Defun(tf.float32, tf.float32, tf.float32,%0A                      func_name=%22attention_sum%22,%0A                      noinline=True)%0A      def attention_sum(v_att, inputs, state):%0A        %22%22%22Calculates a batch- and timweise dot product%22%22%22%0A        return tf.reduce_sum(v_att * (inputs + tf.expand_dims(state, 1)), %5B2%5D)%0A      return attention_sum(v_att, inputs, state)
 %0A%0A  

Simplify dot product