import tensorflow as tf

                                # 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)
            output = tf.reshape(output, [-1, 2 * self.config.hidden_size])

            atten_hidden = tf.tanh(
