[OpAttr]num_rows/num_colums of eye support Tensor type (#45427)

* [OpAttr]num_rows/num_colums of eye support Tensor type

* fix attr cast with long type

paddle/fluid/framework/infershape_utils.cc

@@ -490,6 +490,10 @@ CompatInferMetaContext BuildInferMetaContext(InferShapeContext* ctx,
               infer_meta_context.EmplaceBackAttr(
                   phi::Scalar(PADDLE_GET_CONST(int, attr)));
               break;
+            case framework::proto::AttrType::LONG:
+              infer_meta_context.EmplaceBackAttr(
+                  phi::Scalar(PADDLE_GET_CONST(int64_t, attr)));
+              break;
             case framework::proto::AttrType::STRING:
               infer_meta_context.EmplaceBackAttr(
                   phi::Scalar(PADDLE_GET_CONST(std::string, attr)));

paddle/fluid/framework/operator.cc

@@ -2753,6 +2753,10 @@ void OperatorWithKernel::BuildPhiKernelContext(
               phi_kernel_context->EmplaceBackAttr(std::move(
                   phi::Scalar(PADDLE_GET_CONST(int, attr_iter->second))));
               break;
+            case proto::AttrType::LONG:
+              phi_kernel_context->EmplaceBackAttr(std::move(
+                  phi::Scalar(PADDLE_GET_CONST(int64_t, attr_iter->second))));
+              break;
             case proto::AttrType::STRING:
               phi_kernel_context->EmplaceBackAttr(std::move(phi::Scalar(
                   PADDLE_GET_CONST(std::string, attr_iter->second))));

paddle/fluid/imperative/prepared_operator.h

@@ -420,6 +420,10 @@ void BuildDygraphPhiKernelContext(const phi::KernelSignature& kernel_signature,
               kernel_ctx->EmplaceBackAttr(
                   std::move(phi::Scalar(PADDLE_GET_CONST(int, attr))));
               break;
+            case framework::proto::AttrType::LONG:
+              kernel_ctx->EmplaceBackAttr(
+                  std::move(phi::Scalar(PADDLE_GET_CONST(int64_t, attr))));
+              break;
             case framework::proto::AttrType::STRING:
               kernel_ctx->EmplaceBackAttr(
                   std::move(phi::Scalar(PADDLE_GET_CONST(std::string, attr))));

paddle/fluid/operators/eye_op.cc

@@ -50,11 +50,13 @@ class EyeOpMaker : public framework::OpProtoAndCheckerMaker {
                  "Output data type")
         .SetDefault(framework::proto::VarType::FP32);
     AddAttr<int64_t>("num_rows",
-                     "(int64_t) the number of rows in output tensor");
+                     "(int64_t) the number of rows in output tensor")
+        .SupportTensor();
     AddAttr<int64_t>("num_columns",
                      "(int64_t) the number of columns in output tensor."
                      "Default -1 means that num_columns=num_rows")
-        .SetDefault(-1);
+        .SetDefault(-1)
+        .SupportTensor();
     AddOutput("Out",
               "(Tensor) Construct an identity tensor with "
               "specified shape [num_rows, num_columns]");

paddle/phi/api/yaml/legacy_api.yaml

@@ -874,7 +874,7 @@
   backward : exponential__grad
 
 - api : eye
-  args : (int64_t num_rows, int64_t num_columns, DataType dtype=DataType::FLOAT32, Place place={})
+  args : (Scalar num_rows, Scalar num_columns, DataType dtype=DataType::FLOAT32, Place place={})
   output : Tensor(out)
   infer_meta :
     func : EyeInferMeta

paddle/phi/infermeta/nullary.cc

@@ -51,12 +51,25 @@ void CreateInferMetaBase(const std::vector<int64_t>& shape,
   out->set_layout(layout);
 }
 
-void EyeInferMeta(int64_t num_rows,
-                  int64_t num_columns,
+void EyeInferMeta(const Scalar& num_rows,
+                  const Scalar& num_columns,
                   DataType dtype,
-                  MetaTensor* out) {
-  if (num_columns == -1) num_columns = num_rows;
-  out->set_dims({num_rows, num_columns});
+                  MetaTensor* out,
+                  MetaConfig config) {
+  int64_t rows, columns;
+  if (!config.is_runtime && num_rows.FromTensor()) {
+    rows = -1;
+  } else {
+    rows = num_rows.to<int64_t>();
+  }
+
+  if (!config.is_runtime && num_columns.FromTensor()) {
+    columns = -1;
+  } else {
+    columns = num_columns.to<int64_t>();
+    if (columns == -1) columns = rows;
+  }
+  out->set_dims({rows, columns});
   out->set_dtype(dtype);
 }
 

paddle/phi/infermeta/nullary.h

@@ -15,6 +15,7 @@ limitations under the License. */
 #pragma once
 
 #include "paddle/phi/common/int_array.h"
+#include "paddle/phi/common/scalar.h"
 #include "paddle/phi/core/meta_tensor.h"
 
 namespace phi {
@@ -41,10 +42,11 @@ void CreateInferMetaBase(const std::vector<int64_t>& shape,
                          DataLayout layout,
                          MetaTensor* out);
 
-void EyeInferMeta(int64_t num_rows,
-                  int64_t num_columns,
+void EyeInferMeta(const Scalar& num_rows,
+                  const Scalar& num_columns,
                   DataType dtype,
-                  MetaTensor* out);
+                  MetaTensor* out,
+                  MetaConfig config = MetaConfig());
 
 void GaussianRandomInferMeta(const IntArray& shape,
                              float mean,

paddle/phi/kernels/eye_kernel.h

@@ -14,14 +14,15 @@
 
 #pragma once
 
+#include "paddle/phi/common/scalar.h"
 #include "paddle/phi/core/dense_tensor.h"
 
 namespace phi {
 
 template <typename T, typename Context>
 void EyeKernel(const Context& ctx,
-               int64_t num_rows,
-               int64_t num_columns,
+               const Scalar& num_rows,
+               const Scalar& num_columns,
                DataType dtype,
                DenseTensor* out);
 

paddle/phi/kernels/impl/eye_kernel_impl.h

@@ -15,6 +15,7 @@
 #pragma once
 
 #include "paddle/fluid/platform/for_range.h"
+#include "paddle/phi/common/scalar.h"
 #include "paddle/phi/kernels/funcs/math_function.h"
 
 namespace phi {
@@ -34,20 +35,21 @@ struct EyeFunctor {
 
 template <typename T, typename Context>
 void EyeKernel(const Context& ctx,
-               int64_t num_rows,
-               int64_t num_columns,
+               const Scalar& num_rows,
+               const Scalar& num_columns,
                DataType dtype,
                DenseTensor* out) {
-  auto num = num_columns;
-  if (num == -1) {
-    num = num_rows;
+  auto columns = num_columns.to<int64_t>();
+  auto rows = num_rows.to<int64_t>();
+  if (columns == -1) {
+    columns = rows;
   }
   T* out_data = ctx.template Alloc<T>(out);
   phi::funcs::SetConstant<Context, T> set_zero;
   set_zero(ctx, out, static_cast<T>(0));
-  int64_t num_eyes = (std::min)(num_rows, num);
+  int64_t num_eyes = (std::min)(rows, columns);
   paddle::platform::ForRange<Context> for_range(ctx, num_eyes);
-  EyeFunctor<T> functor(num, out_data);
+  EyeFunctor<T> functor(columns, out_data);
   for_range(functor);
 }
 

python/paddle/fluid/layers/tensor.py

@@ -1791,11 +1791,17 @@ def eye(num_rows,
 
     """
 
+    def _check_attr(attr, message):
+        if isinstance(attr, ((Variable, core.VarBase, core.eager.Tensor))):
+            assert len(attr.shape) == 1 and attr.shape[0] in [1, -1]
+        elif not isinstance(attr, int) or attr < 0:
+            raise TypeError("{} should be a non-negative int.".format(message))
+
+    _check_attr(num_rows, "num_rows")
     if not isinstance(dtype, core.VarDesc.VarType):
         dtype = convert_np_dtype_to_dtype_(dtype)
     if num_columns is not None:
-        if not isinstance(num_columns, int) or num_columns < 0:
-            raise TypeError("num_columns should be a non-negative int")
+        _check_attr(num_columns, "num_columns")
     else:
         num_columns = num_rows
 
@@ -1809,8 +1815,6 @@ def eye(num_rows,
         helper = LayerHelper("eye", **locals())
         check_dtype(dtype, 'dtype',
                     ['float16', 'float32', 'float64', 'int32', 'int64'], 'eye')
-        if not isinstance(num_rows, int) or num_rows < 0:
-            raise TypeError("num_rows should be a non-negative int")
         out = helper.create_variable_for_type_inference(dtype=dtype)
         helper.append_op(type='eye',
                          inputs={},

python/paddle/fluid/tests/unittests/test_eye_op.py

@@ -14,6 +14,7 @@
 
 from __future__ import print_function
 
+import os
 import unittest
 import numpy as np
 from op_test import OpTest
@@ -22,6 +23,9 @@ import paddle
 import paddle.fluid as fluid
 import paddle.fluid.framework as framework
 
+from paddle.fluid.framework import program_guard, Program
+from test_attribute_var import UnittestBase
+
 
 class TestEyeOp(OpTest):
 
@@ -162,5 +166,69 @@ class API_TestTensorEye(unittest.TestCase):
             self.assertRaises(TypeError, test_num_columns_type_check1)
 
 
+class TestEyeRowsCol(UnittestBase):
+
+    def init_info(self):
+        self.shapes = [[2, 3, 4]]
+        self.save_path = os.path.join(self.temp_dir.name, self.path_prefix())
+
+    def test_static(self):
+        main_prog = Program()
+        starup_prog = Program()
+        with program_guard(main_prog, starup_prog):
+            fc = paddle.nn.Linear(4, 10)
+            x = paddle.randn([2, 3, 4])
+            x.stop_gradient = False
+            feat = fc(x)  # [2,3,10]
+
+            tmp = self.call_func(feat)
+            out = feat + tmp
+
+            sgd = paddle.optimizer.SGD()
+            sgd.minimize(paddle.mean(out))
+            self.assertTrue(self.var_prefix() in str(main_prog))
+
+            exe = paddle.static.Executor()
+            exe.run(starup_prog)
+            res = exe.run(fetch_list=[tmp, out])
+            gt = np.eye(3, 10)
+            np.testing.assert_allclose(res[0], gt)
+            paddle.static.save_inference_model(self.save_path, [x], [tmp, out],
+                                               exe)
+            # Test for Inference Predictor
+            infer_outs = self.infer_prog()
+            np.testing.assert_allclose(infer_outs[0], gt)
+
+    def path_prefix(self):
+        return 'eye_rows_cols'
+
+    def var_prefix(self):
+        return "Var["
+
+    def call_func(self, x):
+        rows = paddle.assign(3)
+        cols = paddle.assign(10)
+        out = paddle.eye(rows, cols)
+        return out
+
+    def test_error(self):
+        with self.assertRaises(TypeError):
+            paddle.eye(-1)
+
+
+class TestEyeRowsCol2(TestEyeRowsCol):
+
+    def call_func(self, x):
+        rows = paddle.assign(3)
+        cols = paddle.assign(10)
+        out = paddle.fluid.layers.eye(rows, cols)
+        return out
+
+    def test_error(self):
+        with self.assertRaises(TypeError):
+            paddle.fluid.layers.eye(-1)
+
+
 if __name__ == "__main__":
+    paddle.enable_static()
     unittest.main()

python/paddle/tensor/creation.py

@@ -710,16 +710,20 @@ def eye(num_rows, num_columns=None, dtype=None, name=None):
           #  [0 1 0]]
     """
 
+    def _check_attr(attr, message):
+        if isinstance(attr, ((Variable, core.VarBase, core.eager.Tensor))):
+            assert len(attr.shape) == 1 and attr.shape[0] in [1, -1]
+        elif not isinstance(attr, int) or attr < 0:
+            raise TypeError("{} should be a non-negative int.".format(message))
+
+    _check_attr(num_rows, "num_rows")
+
     if dtype is None:
         dtype = 'float32'
-    if num_columns is None:
-        num_columns = num_rows
-
     if not isinstance(dtype, core.VarDesc.VarType):
         dtype = convert_np_dtype_to_dtype_(dtype)
     if num_columns is not None:
-        if not isinstance(num_columns, int) or num_columns < 0:
-            raise TypeError("num_columns should be a non-negative int")
+        _check_attr(num_columns, "num_columns")
     else:
         num_columns = num_rows
 
@@ -735,8 +739,6 @@ def eye(num_rows, num_columns=None, dtype=None, name=None):
         helper = LayerHelper("eye", **locals())
         check_dtype(dtype, 'dtype',
                     ['float16', 'float32', 'float64', 'int32', 'int64'], 'eye')
-        if not isinstance(num_rows, int) or num_rows < 0:
-            raise TypeError("num_rows should be a non-negative int")
         out = helper.create_variable_for_type_inference(dtype=dtype)
         helper.append_op(type='eye',
                          inputs={},