Make UnitNormalization layer stateless.

There is no need to resolve negative axes in `build`, as `tf.linalg.l2_normalize` can handle them. Kept the build method to validate the axes in the context of the `input_shape`. Also added call to `super.build(...)` per best practice on Keras 2.

Note that in Keras 3, `UnitNormalization` is already stateless.

PiperOrigin-RevId: 699253410

Author: hertschuh

tf_keras/layers/normalization/unit_normalization.py

@@ -60,7 +60,8 @@ def __init__(self, axis=-1, **kwargs):
         self.supports_masking = True
 
     def build(self, input_shape):
-        self.axis = tf_utils.validate_axis(self.axis, input_shape)
+        tf_utils.validate_axis(self.axis, input_shape)
+        super().build(input_shape)
 
     def call(self, inputs):
         inputs = tf.cast(inputs, self.compute_dtype)