Below is the same model implemented by Tensorflow.
# hyper parameters
batch_size = 8
word_dimension = 300
pos_dimension = 40
gru_hidden_dimension = 200
score_hidden_dimension = 2 * gru_hidden_dimension

learning_rate = 0.001
beta = 1e-4

# initialize variables
train_set = Reader(train_path)
valid_set = Reader(valid_path)
test_set = Reader(test_path)

with tf.device("/gpu:0"):
    if is_use_embedding:
        generator = EmbeddingGenerator(train_path, valid_path, test_path)
        word_dict = generator.get_word_dict()
        _, pos_dict, arc_dict = train_set.get_dicts()
        word_embedding_generated = generator.get_word_embeddings(embedding_path, word_dimension)
        word_embedding = tf.Variable(name='word_embedding',
                                     initial_value=word_embedding_generated,
                                     dtype=tf.float32)
    else:
        word_dict, pos_dict, arc_dict = train_set.get_dicts()
        word_embedding = tf.Variable(name='word_embedding',
                                     initial_value=tf.random_uniform([len(word_dict), word_dimension], -0.1, 0.1),
                                     dtype=tf.float32)

    pos_embedding = tf.Variable(name='pos_embedding',
                                initial_value=tf.random_uniform([len(pos_dict), pos_dimension], -0.1, 0.1),
                                dtype=tf.float32)

    # placeholders
    batch_max_length = tf.placeholder(name='batch_max_length', dtype=tf.int32, shape=[])
    length = tf.placeholder(name='length', dtype=tf.int32, shape=[batch_size])

    word_indices = tf.placeholder(name='words_indices', shape=[batch_size, None], dtype=tf.int32)
    pos_indices = tf.placeholder(name='pos_indices', shape=[batch_size, None], dtype=tf.int32)

    index_matrix = tf.placeholder(name='index_matrix', shape=[batch_size, None, None, 2], dtype=tf.int32)
    gold_score_matrix = tf.placeholder(name='gold_score_matrix', shape=[batch_size, None, None], dtype=tf.float32)

    # build the network
    words = tf.nn.embedding_lookup(name='words', params=word_embedding, ids=word_indices)
    pos = tf.nn.embedding_lookup(name='pos', params=pos_embedding, ids=pos_indices)

    x = tf.concat(name='words_c_pos', values=[words, pos], axis=2)

    gru_fw = rnn.GRUCell(num_units=gru_hidden_dimension)
    gru_bw = rnn.GRUCell(num_units=gru_hidden_dimension)

    init_fw = gru_fw.zero_state(batch_size=batch_size, dtype=tf.float32)
    init_bw = gru_bw.zero_state(batch_size=batch_size, dtype=tf.float32)

    with tf.variable_scope('Bi-GRU'):
        outputs, states = tf.nn.bidirectional_dynamic_rnn(cell_fw=gru_fw,
                                                          cell_bw=gru_bw,
                                                          inputs=x,
                                                          sequence_length=length,
                                                          initial_state_fw=init_fw,
                                                          initial_state_bw=init_bw)

    # a bit different from Zhang et al. 2016
    c = tf.concat(outputs, 2)

    loss = tf.get_variable(name='loss', dtype=tf.float32, shape=[])

    u = tf.get_variable(name='u',
                        shape=[4 * gru_hidden_dimension, 2 * gru_hidden_dimension],
                        dtype=tf.float32,
                        initializer=tf.random_uniform_initializer(-0.1, 0.1))

    v = tf.get_variable(name='v',
                        shape=[2 * gru_hidden_dimension, 1],
                        dtype=tf.float32,
                        initializer=tf.random_uniform_initializer(-0.1, 0.1))

    g_s = []

    for i in range(batch_size):

        actual_length = length[i]
        sentence_rep = c[i]
        actual_gold_score = gold_score_matrix[i]

        rep_matrix = tf.reshape(tensor=tf.nn.embedding_lookup(params=sentence_rep, ids=index_matrix[i]),
                                shape=[batch_max_length, batch_max_length, 4 * gru_hidden_dimension],
                                name='rep_matrix')

        h_score = tf.nn.tanh(tf.einsum('ijk,kl->ijl', rep_matrix, u))

        o_score = tf.squeeze(tf.einsum('ijk,kl->ijl', h_score, v), 2)

        o_cut = o_score[:actual_length, :actual_length]
        o_rm = o_cut[1:, :]
        gs_cut = actual_gold_score[:actual_length, :actual_length]
        gs_rm = gs_cut[1:, :]

        g = tf.clip_by_value(tf.nn.softmax(o_rm), 1e-10, 1e10)
        g_s.append(g)

        loss += -tf.reduce_sum(gs_rm * tf.log(g))

    loss += tf.add_n([tf.nn.l2_loss(v) for v in tf.trainable_variables()]) * beta

    opt = tf.train.AdamOptimizer(learning_rate).minimize(loss)