def rnn_fast_weights(seqs,
                     batch_size=32,
                     hidden_size=50,
                     fast_steps=1,
                     fast_decay_rate=.95,
                     fast_lr=.5):
  """Implements a RNN with fast weights.

  The RNN updates a fast-weights "memory" matrix A via Hebbian-style updates
  as described by [Ba et al. 2018] Using Fast Weights To Attend To the Recent
  Past (https://arxiv.org/abs/1610.062580).

  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 in the RNN.
    fast_steps (int): Number of inner loop iterations we apply fast weights.
    fast_decay_rate (float): Decay rate (lambda) for fast weights update.
    fast_lr (float): Learning rate (eta) for fast weights update.

  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]
  state = tf.zeros([batch_size, hidden_size], name="state")
  # a_memory: <tf.float32>[batch_size, hidden_size, hidden_size]
  a_memory = tf.zeros(
      [batch_size, hidden_size, hidden_size],
      dtype=tf.float32,
      name="A")

  # Unroll graph manually
  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}).
    inp = tf.concat([seqs[:, timestep, :], state], 1)
    state = tf.matmul(inp, w1) + b1
    boundary_state = state  # "sustained boundary condition" pre-nonlinearity
    state = tf.tanh(state)

    # Apply fast weights
    for _ in range(fast_steps):
      # fw_state: <tf.float32>[batch_size, hidden_size]
      fw_state = tf.squeeze(tf.matmul(a_memory, tf.expand_dims(state, 2)))
      state = boundary_state + fw_state
      state = layer_norm(state, hidden_size)
      state = tf.tanh(state)

    # Update fast weights matrix
    a_memory = fast_decay_rate * a_memory + fast_lr * tf.matmul(
        tf.expand_dims(state, 2),  # <tf.float32>[batch_size, hidden_size, 1]
        tf.expand_dims(state, 1))  # <tf.float32>[batch_size, 1, hidden_size]

  # preds: <tf.float32>[batch_size, pattern_size]
  preds = tf.layers.dense(state, pattern_size)

  return preds
Implements a RNN with fast weights.

  The RNN updates a fast-weights "memory" matrix A via Hebbian-style updates
  as described by [Ba et al. 2018] Using Fast Weights To Attend To the Recent
  Past (https://arxiv.org/abs/1610.062580).

  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 in the RNN.
    fast_steps (int): Number of inner loop iterations we apply fast weights.
    fast_decay_rate (float): Decay rate (lambda) for fast weights update.
    fast_lr (float): Learning rate (eta) for fast weights update.

  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