Update examples/keras_rs/stu.py

Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com>

Author: LakshmiKalaKadali

examples/keras_rs/stu.py

@@ -317,10 +317,7 @@ def keras_concat_2D_jagged_resolver(
 
 def keras_layer_norm(x, weight, bias, eps):
     # Functional Layer Norm steps
-    mean = ops.mean(x, axis=-1, keepdims=True)
-    variance = ops.mean(ops.square(x - mean), axis=-1, keepdims=True)
-    std = ops.sqrt(variance + eps)
-    normalized_x = (x - mean) / std
+    normalized_x = ops.layer_norm(x, axis=-1, epsilon=eps)
     return normalized_x * weight + bias