improve rmsln preformance when torch backend

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):
@@ -2725,6 +2726,10 @@ def compute_output_spec(self, x):
         return KerasTensor(shape=x.shape)
 
     def call(self, x):
+        if backend.backend() == "torch" and is_continuous_axis(self.axis):
+            import torch.nn.functional as F
+
+            return F.rms_norm(x, self.axis, self.scale, self.epsilon)
         return _rms_normalization(
             x, scale=self.scale, axis=self.axis, epsilon=self.epsilon
         )

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 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