output layer code does not match paper

[yuzhoul1991]

In the paper it gives the impression that when M is computed, you feed that to a bidirectional LSTM to get M^2. However in the code base there are a bunch of intermediate step such as:

    a1i = softsel(tf.reshape(g1, [N, M * JX, 2 * d]), tf.reshape(logits, [N, M * JX]))
        a1i = tf.tile(tf.expand_dims(tf.expand_dims(a1i, 1), 1), [1, M, JX, 1])

        (fw_g2, bw_g2), _ = bidirectional_dynamic_rnn(d_cell4_fw, d_cell4_bw, tf.concat(axis=3, values=[p0, g1, a1i, g1 * a1i]),
                                                      x_len, dtype='float', scope='g2')  # [N, M, JX, 2d]

Can you explain what these steps are doing? Why does the paper not stating these steps? Did i miss something?