[NewIR] No.7 Migrate paddle.softmaxinto pir (#57415)

Difers · web-flow · commit 571f32c7199d · 2023-09-22T14:31:08.000+08:00
diff --git a/python/paddle/nn/functional/activation.py b/python/paddle/nn/functional/activation.py
@@ -1196,9 +1196,13 @@ def softmax(x, axis=-1, dtype=None, name=None):
               [0.03205860, 0.08714432, 0.23688282, 0.64391426]]])
     """
 
-    if (dtype is not None) and (not isinstance(dtype, core.VarDesc.VarType)):
+    if (
+        (dtype is not None)
+        and (not isinstance(dtype, core.VarDesc.VarType))
+        and (not isinstance(dtype, core.DataType))
+    ):
         dtype = convert_np_dtype_to_dtype_(dtype)
-    if in_dynamic_mode():
+    if in_dynamic_or_pir_mode():
         outs_cast = x if dtype is None else _C_ops.cast(x, dtype)
         return _C_ops.softmax(outs_cast, axis)
     else:
diff --git a/test/legacy_test/test_softmax_op.py b/test/legacy_test/test_softmax_op.py
@@ -84,9 +84,9 @@ def test_check_output(self):
         # TODO(wangzhongpu): support mkldnn op in dygraph mode
         if self.use_cudnn:
             place = core.CUDAPlace(0)
-            self.check_output_with_place(place, atol=1e-5)
+            self.check_output_with_place(place, atol=1e-5, check_new_ir=True)
         else:
-            self.check_output(check_prim=True)
+            self.check_output(check_prim=True, check_new_ir=True)
 
     def test_check_grad(self):
         # TODO(wangzhongpu): support mkldnn op in dygraph mode
@@ -99,6 +99,7 @@ def test_check_grad(self):
                     "Out",
                     max_relative_error=0.01,
                     check_dygraph=(not self.use_mkldnn),
+                    check_new_ir=True,
                 )
         else:
             self.check_grad(
@@ -107,6 +108,7 @@ def test_check_grad(self):
                 max_relative_error=0.01,
                 check_dygraph=(not self.use_mkldnn),
                 check_prim=True,
+                check_new_ir=True,
             )
 
 
@@ -144,9 +146,9 @@ def test_check_output(self):
         # TODO(wangzhongpu): support mkldnn op in dygraph mode
         if self.use_cudnn:
             place = core.CUDAPlace(0)
-            self.check_output_with_place(place, atol=1e-5)
+            self.check_output_with_place(place, atol=1e-5, check_new_ir=True)
         else:
-            self.check_output(check_prim=True)
+            self.check_output(check_prim=True, check_new_ir=True)
 
 
 @unittest.skipIf(
@@ -178,9 +180,9 @@ def test_check_output(self):
         # TODO(wangzhongpu): support mkldnn op in dygraph mode
         if self.use_cudnn:
             place = core.CUDAPlace(0)
-            self.check_output_with_place(place, atol=1e-5)
+            self.check_output_with_place(place, atol=1e-5, check_new_ir=True)
         else:
-            self.check_output(check_prim=True)
+            self.check_output(check_prim=True, check_new_ir=True)
 
 
 class TestSoftmaxOp2(TestSoftmaxOp):
diff --git a/test/mkldnn/test_softmax_mkldnn_op.py b/test/mkldnn/test_softmax_mkldnn_op.py
@@ -27,8 +27,11 @@
     TestSoftmaxOp_ZeroDim1,
 )
 
+import paddle
 from paddle.base import core
 
+paddle.enable_static()
+
 
 def stable_softmax(x):
     """Compute the softmax of vector x in a numerically stable way."""

Original file line number	Diff line number	Diff line change
`@@ -27,8 +27,11 @@`
`27`	`27`	`TestSoftmaxOp_ZeroDim1,`
`28`	`28`	`)`
`29`	`29`
	`30`	`+import paddle`
`30`	`31`	`from paddle.base import core`
`31`	`32`
	`33`	`+paddle.enable_static()`
	`34`	`+`
`32`	`35`
`33`	`36`	`def stable_softmax(x):`
`34`	`37`	`"""Compute the softmax of vector x in a numerically stable way."""`