def _init_norm(self):
        """Set the norm of the weight vector."""
        kernel_norm = tf.sqrt(tf.reduce_sum(tf.square(self.weight_v), axis=self.kernel_norm_axes))
        self.weight_g.assign(kernel_norm[:, tf.newaxis, tf.newaxis])
Set the norm of the weight vector.
dpkd/transformers/src/transformers/models/wav2vec2/modeling_tf_wav2vec2.py