keras/src/layers/normalization/unit_normalization.py/UnitNormalization/call
class UnitNormalization:    def call(self, inputs):
        x = ops.cast(inputs, self.compute_dtype)

        square_sum = ops.sum(ops.square(x), axis=self.axis, keepdims=True)
        x_inv_norm = ops.rsqrt(ops.maximum(square_sum, 1e-12))
        return ops.multiply(x, x_inv_norm)