[phi] fix sequence_mask dtype

paddle/fluid/pir/dialect/operator/ir/update_ops.yaml

@@ -17,7 +17,7 @@
   support_tensor : [start, end, step]
 
 - op : sequence_mask
-  args: (Tensor x, Scalar(int) max_len, int out_dtype)
+  args: (Tensor x, Scalar(int) max_len, DataType out_dtype)
   output: Tensor(y)
   infer_meta:
     func: SequenceMaskScalarInferMeta

paddle/phi/api/yaml/legacy_ops.yaml

@@ -981,7 +981,7 @@
   backward : rrelu_grad
 
 - op : sequence_mask
-  args: (Tensor x, Scalar(int) max_len, int out_dtype)
+  args: (Tensor x, Scalar(int) max_len, DataType out_dtype)
   output: Tensor(y)
   infer_meta:
     func: SequenceMaskScalarInferMeta

paddle/phi/infermeta/binary.cc

@@ -2835,7 +2835,7 @@ void SearchsortedInferMeta(const MetaTensor& sorted_sequence,
 void SequenceMaskInferMeta(const MetaTensor& x,
                            const MetaTensor& max_len_tensor,
                            int maxlen,
-                           int out_dtype,
+                           DataType out_dtype,
                            MetaTensor* y) {
   auto dim = common::vectorize<int>(x.dims());
 
@@ -2846,8 +2846,7 @@ void SequenceMaskInferMeta(const MetaTensor& x,
   }
 
   y->set_dims(common::make_ddim(dim));
-  auto out_phi_dtype = phi::TransToPhiDataType(out_dtype);
-  y->set_dtype(out_phi_dtype);
+  y->set_dtype(out_dtype);
 }
 
 void SoftmaxMaskFuseInferMeta(const MetaTensor& x,

paddle/phi/infermeta/binary.h

@@ -14,6 +14,7 @@ limitations under the License. */
 
 #pragma once
 
+#include "paddle/phi/common/data_type.h"
 #include "paddle/phi/common/int_array.h"
 #include "paddle/phi/common/scalar.h"
 #include "paddle/phi/core/meta_tensor.h"
@@ -452,7 +453,7 @@ void SearchsortedInferMeta(const MetaTensor& sorted_sequence,
 void SequenceMaskInferMeta(const MetaTensor& x,
                            const MetaTensor& max_len_tensor,
                            int maxlen,
-                           int out_dtype,
+                           DataType out_dtype,
                            MetaTensor* y);
 
 void SoftmaxMaskFuseInferMeta(const MetaTensor& x,

paddle/phi/infermeta/unary.cc

@@ -4018,14 +4018,13 @@ void SplitWithNumInferMeta(const MetaTensor& x,
 
 void SequenceMaskScalarInferMeta(const MetaTensor& x,
                                  const Scalar& max_len,
-                                 int out_dtype,
+                                 DataType out_dtype,
                                  MetaTensor* y) {
   auto dim = phi::vectorize<int>(x.dims());
   int maxlen = max_len.to<int>();
   dim.push_back(maxlen > 0 ? maxlen : -1);
   y->set_dims(phi::make_ddim(dim));
-  auto out_phi_dtype = phi::TransToPhiDataType(out_dtype);
-  y->set_dtype(out_phi_dtype);
+  y->set_dtype(out_dtype);
 }
 
 void SquaredL2NormInferMeta(const MetaTensor& x, MetaTensor* out) {

paddle/phi/infermeta/unary.h

@@ -555,7 +555,7 @@ void RReluGradInferMeta(const MetaTensor& out_grad,
 
 void SequenceMaskScalarInferMeta(const MetaTensor& x,
                                  const Scalar& max_len,
-                                 int out_dtype,
+                                 DataType out_dtype,
                                  MetaTensor* y);
 
 void SetValueInferMeta(const MetaTensor& x, MetaTensor* out);

paddle/phi/kernels/impl/sequence_mask_kernel_impl.h

@@ -31,7 +31,7 @@ template <typename T, typename Context>
 void SequenceMaskScalarKernel(const Context& ctx,
                               const DenseTensor& x,
                               const Scalar& max_len,
-                              int out_dtype,
+                              DataType out_dtype,
                               DenseTensor* y) {
   int maxlen = max_len.to<int>();
   auto* x_data = x.data<T>();
@@ -58,7 +58,7 @@ void SequenceMaskScalarKernel(const Context& ctx,
     y->Resize(common::make_ddim(y_dim));
   }
 
-  phi::VisitDataType(phi::TransToPhiDataType(out_dtype),
+  phi::VisitDataType(out_dtype,
                      phi::funcs::SequenceMaskFunctor<Context, T>(
                          ctx, x_data, y, x_numel * maxlen, maxlen));
 }
@@ -68,7 +68,7 @@ void SequenceMaskKernel(const Context& ctx,
                         const DenseTensor& x,
                         const paddle::optional<DenseTensor>& max_len_tensor,
                         int maxlen,
-                        int out_dtype,
+                        DataType out_dtype,
                         DenseTensor* y) {
   if (max_len_tensor) {
     bool is_gpu_place = ctx.GetPlace().GetType() == phi::AllocationType::GPU;

paddle/phi/kernels/sequence_mask_kernel.h

@@ -19,7 +19,7 @@ void SequenceMaskKernel(const Context& ctx,
                         const DenseTensor& x,
                         const paddle::optional<DenseTensor>& max_len_tensor,
                         int maxlen,
-                        int out_dtype,
+                        DataType out_dtype,
                         DenseTensor* y);
 
 }  // namespace phi

python/paddle/nn/functional/extension.py

@@ -16,6 +16,7 @@
 
 
 from paddle import _C_ops, tensor
+from paddle.pir.core import vartype_to_datatype
 from paddle.utils import deprecated
 
 from ...base.data_feeder import check_type, check_variable_and_dtype
@@ -25,6 +26,7 @@
     convert_np_dtype_to_dtype_,
     core,
     in_dynamic_or_pir_mode,
+    in_pir_mode,
 )
 
 __all__ = []
@@ -103,10 +105,13 @@ def sequence_mask(x, maxlen=None, dtype='int64', name=None):
     if in_dynamic_or_pir_mode():
         if not isinstance(dtype, (core.VarDesc.VarType, core.DataType)):
             dtype = convert_np_dtype_to_dtype_(dtype)
-        if maxlen is not None:
-            out = _C_ops.sequence_mask(x, maxlen, dtype)
-            out.stop_gradient = True
-            return out
+        if maxlen is None:
+            maxlen = -1
+        if in_pir_mode() and isinstance(dtype, core.VarDesc.VarType):
+            dtype = vartype_to_datatype[dtype]
+        out = _C_ops.sequence_mask(x, maxlen, dtype)
+        out.stop_gradient = True
+        return out
 
     helper = LayerHelper('sequence_mask', **locals())
     out = helper.create_variable_for_type_inference(dtype=dtype)