update

Author: pass-lin

keras/src/ops/nn.py

@@ -2726,10 +2726,6 @@ 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
         )
@@ -2777,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

@@ -7,7 +7,7 @@
 
 def is_continuous_axis(axis):
     # Used to determine whether the dimensions in an axis are continuous
-    if len(axis) == 1:
+    if isinstance(axis, int) or len(axis) == 1:
         return True
     positive_order_flag = True
     for i in range(len(axis) - 1):