improve RMSNorm preformance when torch backend (#21325)

* improve rmsln preformance when torch backend

* update

Author: pass-lin

keras/src/ops/nn.py

@@ -14,6 +14,7 @@
 from keras.src.ops import operation_utils
 from keras.src.ops.operation import Operation
 from keras.src.ops.operation_utils import reduce_shape
+from keras.src.utils.python_utils import is_continuous_axis
 
 
 class Relu(Operation):
@@ -2772,6 +2773,14 @@ def rms_normalization(x, scale=1, axis=-1, epsilon=None):
 
 
 def _rms_normalization(x, scale=1, axis=-1, epsilon=None):
+    if backend.backend() == "torch" and is_continuous_axis(axis):
+        import torch.nn.functional as F
+
+        if isinstance(axis, (tuple, list)):
+            normalized_shape = tuple([x.shape[dim] for dim in axis])
+        else:
+            normalized_shape = x.shape[axis]
+        return F.rms_norm(x, normalized_shape, scale, epsilon)
     x = backend.convert_to_tensor(x)
     if len(x.shape) == 0:
         x = backend.numpy.expand_dims(x, axis=0)

keras/src/utils/python_utils.py

@@ -5,6 +5,24 @@
 import types as python_types
 
 
+def is_continuous_axis(axis):
+    # Used to determine whether the dimensions in an axis are continuous
+    if isinstance(axis, int) or len(axis) == 1:
+        return True
+    positive_order_flag = True
+    for i in range(len(axis) - 1):
+        if axis[i + 1] - axis[i] != 1:
+            positive_order_flag = False
+            break
+
+    negative_order_flag = True
+    for i in range(len(axis) - 1):
+        if axis[i + 1] - axis[i] != 1:
+            negative_order_flag = False
+            break
+    return positive_order_flag or negative_order_flag
+
+
 def default(method):
     """Decorates a method to detect overrides in subclasses."""
     method._is_default = True