PaddlePaddle · SigureMo · Jan 17, 2024 · Dec 28, 2023 · Dec 28, 2023 · Dec 28, 2023
diff --git a/paddle/fluid/pir/dialect/operator/ir/update_ops.yaml b/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

diff --git a/paddle/phi/api/yaml/legacy_ops.yaml b/paddle/phi/api/yaml/legacy_ops.yaml
@@ -991,7 +991,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

diff --git a/paddle/phi/infermeta/binary.cc b/paddle/phi/infermeta/binary.cc
@@ -2857,7 +2857,7 @@ void ShuffleBatchInferMeta(const MetaTensor& x,
 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());
 
@@ -2868,8 +2868,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,

diff --git a/paddle/phi/infermeta/binary.h b/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"
@@ -457,7 +458,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 ShuffleBatchInferMeta(const MetaTensor& x,

diff --git a/paddle/phi/infermeta/unary.cc b/paddle/phi/infermeta/unary.cc
@@ -4097,14 +4097,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) {

diff --git a/paddle/phi/infermeta/unary.h b/paddle/phi/infermeta/unary.h
@@ -581,7 +581,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);

diff --git a/paddle/phi/kernels/impl/sequence_mask_kernel_impl.h b/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;

diff --git a/paddle/phi/kernels/sequence_mask_kernel.h b/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
@@ -235,11 +235,6 @@ def __setattr__(self, name, val):
 }
 
 
-# FIXME(dev): We haven't fully verified eager mode on XPU et.al but
-# only GPU/CPU. Remove this after we improve this feature.
-_is_first_import_ = True
-
-
 def in_dygraph_mode():
     """
 

diff --git a/python/paddle/nn/functional/extension.py b/python/paddle/nn/functional/extension.py
@@ -103,10 +103,11 @@ 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
+        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)

diff --git a/python/paddle/pir/core.py b/python/paddle/pir/core.py
@@ -45,21 +45,6 @@
     VarDesc.VarType.COMPLEX128: DataType.COMPLEX128,
 }
 
-datatype_to_vartype = {
-    DataType.FLOAT32: VarDesc.VarType.FP32,
-    DataType.FLOAT64: VarDesc.VarType.FP64,
-    DataType.FLOAT16: VarDesc.VarType.FP16,
-    DataType.BFLOAT16: VarDesc.VarType.BF16,
-    DataType.INT32: VarDesc.VarType.INT32,
-    DataType.INT16: VarDesc.VarType.INT16,
-    DataType.INT64: VarDesc.VarType.INT64,
-    DataType.BOOL: VarDesc.VarType.BOOL,
-    DataType.UINT8: VarDesc.VarType.UINT8,
-    DataType.INT8: VarDesc.VarType.INT8,
-    DataType.COMPLEX64: VarDesc.VarType.COMPLEX64,
-    DataType.COMPLEX128: VarDesc.VarType.COMPLEX128,
-}
-
 np_type_to_paddle_type = {
     np.dtype("float32"): DataType.FLOAT32,
     np.dtype("float64"): DataType.FLOAT64,

diff --git a/test/legacy_test/prim_op_test.py b/test/legacy_test/prim_op_test.py
@@ -32,12 +32,13 @@
     canonicalize_attrs,
     in_dygraph_mode,
     in_pir_mode,
+    paddle_type_to_proto_type,
     use_pir_api,
 )
 from paddle.decomposition import decompose
 from paddle.incubate.autograd import primapi
 from paddle.jit.dy2static.utils import parse_arg_and_kwargs
-from paddle.pir.core import datatype_to_vartype, vartype_to_datatype
+from paddle.pir.core import vartype_to_datatype
 
 
 def flatten(nest_list):
@@ -156,7 +157,7 @@ def convert_dtype(dtype, target_dtype):
                 isinstance(dtype, paddle.pir.core.DataType)
                 and target_dtype is core.VarDesc.VarType
             ):
-                return datatype_to_vartype[dtype]
+                return paddle_type_to_proto_type[dtype]
             return dtype
 
         # NOTE(xiongkun): the logic of constructing parameters: