def rnn_attention(seqs, batch_size=32, hidden_size=50):
  """A basic RNN with attention.

  Args:
    seqs: <tf.float32>[batch_size, total_len, pattern_size] Sequences of
      patterns, where `total_len` is the length of a sequence (including the
      query pattern) and `pattern_size` is the dimensionality of each pattern.
    batch_size (int): Batch size.
    hidden_size (int): Number of hidden units.

  Returns:
    preds: <tf.float32>[batch_size, pattern_size] The retrieved pattern for
      the degraded query at the end of the sequence.
  """
  _, total_len, pattern_size = seqs.shape

  # Initialize parameters
  w1 = tf.get_variable("w1", shape=[pattern_size + hidden_size, hidden_size])
  b1 = tf.get_variable(
      "b1", shape=[hidden_size], initializer=tf.zeros_initializer())

  # state: <tf.float32>[batch_size, hidden_size]
  # Store all intermediate hidden states h_i in all_states, where we will later
  # use for computing the attention-based hidden representation c, where
  # c = \sum_i \alpha_i h_i
  state = tf.zeros([batch_size, hidden_size], name="state")
  all_states = []

  # Construct the unrolled RNN manually for the sequence
  for timestep in range(total_len):
    # concat_state: <tf.float32>[batch_size, pattern_size + hidden_size]
    # For each batch element, take the input for a given timestep and append the
    # hidden state vector. The weight matrix w1 and bias b1 encodes both the W
    # and U matrices in the h_t = g(Wx_i, Uh_{t-1}).
    concat_state = tf.concat([seqs[:, timestep, :], state], 1)
    state = tf.matmul(concat_state, w1) + b1
    state = tf.tanh(state)
    all_states.append(state)

  # Recall that the sequence contains all the input patterns plus a query
  # pattern and that the query patterns is a corrupted form of one of the input
  # patterns we want to recover. Remove last TIMESTEPS_PER_PATTERN=1 states for
  # the query pattern as that's irrelevant for computing the attention-weighted
  # hidden representation.
  #
  # all_states: <tf.float32>[total_len - 1, batch_size, hidden_size].
  all_states = tf.convert_to_tensor(all_states[:-1], dtype=tf.float32)

  # Transpose so batch_size is first dim
  # new shape <tf.float32>[batch_size, total_len - 1, hidden_size]
  all_states = tf.transpose(all_states, [1, 0, 2])

  # Compute similarity between final state and all previous states.
  # attention: <tf.float32>[batch_size, total_len - 1, 1]
  attention = tf.reduce_sum(
      tf.multiply(tf.expand_dims(state, 1), all_states), 2, keep_dims=True)

  # Turn into probability distribution
  # attention_dist: <tf.float32>[batch_size, total_len - 1, 1]
  attention_dist = tf.nn.softmax(attention, dim=1)
  # context: <tf.float32>[batch_size, hidden_size]
  context = tf.reduce_sum(tf.multiply(attention_dist, all_states), 1)

  # concat_state: <tf.float32>[batch_size, 2 * hidden_size]
  concat_state = tf.concat([context, state], 1)
  # preds: <tf.float32>[batch_size, pattern_size]
  preds = tf.layers.dense(concat_state, pattern_size)

  return preds
A basic RNN with attention.

  Args:
    seqs: <tf.float32>[batch_size, total_len, pattern_size] Sequences of
      patterns, where `total_len` is the length of a sequence (including the
      query pattern) and `pattern_size` is the dimensionality of each pattern.
    batch_size (int): Batch size.
    hidden_size (int): Number of hidden units.

  Returns:
    preds: <tf.float32>[batch_size, pattern_size] The retrieved pattern for
      the degraded query at the end of the sequence.
  
language/labs/memory/baseline_models.py