import tensorflow as tf

            V = tf.get_variable(name="attn_V", shape=[2 * self.config.hidden_size, 1],
                                initializer=tf.contrib.layers.xavier_initializer(),
                                # initializer=tf.truncated_normal_initializer(),
                                # initializer=tf.keras.initializers.lecun_normal(),
                                dtype=tf.float32)
            U = tf.get_variable(name="attn_U",
                                shape=[2 * self.config.hidden_size, 2 * self.config.hidden_size],
                                initializer=tf.contrib.layers.xavier_initializer(),
                                # initializer=tf.truncated_normal_initializer(),
                                # initializer=tf.keras.initializers.lecun_normal(),
                                dtype=tf.float32)

            self.position_emb = tf.reshape(self.position_emb, [-1, 2 * self.config.hidden_size])
            shape = tf.shape(output)
