import tensorflow as tf

from ..util import static_shape, static_rank
from ..linear import linear
from ..variable import variable



def embeddings_to_embedding(child_embeddings,
                            output_embedding_size,
                            context_vector_size):
  assert static_rank(child_embeddings) == 3

  with tf.name_scope("embeddings_to_embedding"):
    embedding_size = static_shape(child_embeddings)[2]
    rnn_cell = tf.nn.rnn_cell.GRUCell(output_embedding_size, embedding_size)

    cell_outputs = []
    cell_state = rnn_cell.zero_state(tf.shape(child_embeddings)[0], tf.float32)
    for iteration, child_embedding \
        in enumerate(_split_child_embeddings(child_embeddings)):
      if iteration != 0: tf.get_variable_scope().reuse_variables()
      cell_output, cell_state = rnn_cell(child_embedding, cell_state)
      cell_outputs.append(cell_output)

    return _attention_please(tf.pack(cell_outputs), context_vector_size)


def _attention_please(xs, context_vector_size):
  sequence_length = static_shape(xs)[0]
  batch_size = tf.shape(xs)[1]
  embedding_size = static_shape(xs)[2]

  context_vector = variable([context_vector_size, 1])

  attention = tf.nn.softmax(tf.transpose(tf.reshape(
      tf.matmul(
        tf.tanh(linear(
          tf.reshape(
            xs,
            [-1, embedding_size]), # -1 denotes sequence_length * batch_size
          context_vector_size)),
        context_vector),
      [sequence_length, -1]))) # -1 denotes batch_size

  return tf.transpose(tf.pack([
      tf.squeeze(tf.batch_matmul(tf.expand_dims(x, 1),
                                 tf.expand_dims(attention, 2)),
                 [1, 2])
      for x in tf.unpack(tf.transpose(xs))]))


def _split_child_embeddings(child_embeddings):
  return tf.unpack(tf.transpose(
      child_embeddings,
      [1, 0] + _dimension_indices(child_embeddings)[2:]))


def _dimension_indices(tensor):
  return list(range(static_rank(tensor)))

@@ -1524,22 +1524,152 @@
 f.pack(%5B
-%0A     
+_inner_product(x, attention)%0A                               for x in tf.unpack(tf.transpose(xs))%5D))%0A%0A%0Adef _inner_product(x, y):%0A  return
  tf.sque
@@ -1743,16 +1743,19 @@
         
+   
 tf.expan
@@ -1761,25 +1761,17 @@
 nd_dims(
-attention
+y
 , 2)),%0A 
@@ -1790,16 +1790,19 @@
     
+ 
+  
 %5B1, 2%5D)%0A
     
@@ -1801,54 +1801,8 @@
 2%5D)%0A
-      for x in tf.unpack(tf.transpose(xs))%5D))%0A
 %0A%0Ade

Refactor nn.embedding.embeddings_to_embedding