[Fix] Fix det infermeta and support 0-D tensor for paddle.eye (#68341)

HydrogenSulfate · web-flow · commit 610c3c74490b · 2024-09-23T11:58:53.000+08:00
* correct infer shape in det infermeta

* fix 0D tensor input for paddle.eye

* add unitest

* fix
diff --git a/paddle/phi/infermeta/unary.cc b/paddle/phi/infermeta/unary.cc
@@ -4717,6 +4717,17 @@ void SumInferMeta(const MetaTensor& x,
   SumRawInferMeta(x, axis, keep_dim, reduce_all, dtype, out, config);
 }
 
+void DetInferMeta(const MetaTensor& x, MetaTensor* out, MetaConfig config) {
+  // remove the last two demension
+  auto out_dim = common::vectorize<int>(x.dims());
+  out_dim.pop_back();
+  out_dim.pop_back();
+
+  out->set_dims(common::make_ddim(out_dim));
+  out->set_dtype(x.dtype());
+  out->set_layout(x.layout());
+}
+
 void PartialSumInferMeta(const std::vector<const MetaTensor*>& xs,
                          int start_index,
                          int length,
diff --git a/paddle/phi/infermeta/unary.h b/paddle/phi/infermeta/unary.h
@@ -775,6 +775,10 @@ void SumInferMeta(const MetaTensor& x,
                   MetaTensor* out,
                   MetaConfig config = MetaConfig());
 
+void DetInferMeta(const MetaTensor& x,
+                  MetaTensor* out,
+                  MetaConfig config = MetaConfig());
+
 void SumRawInferMeta(const MetaTensor& x,
                      const IntArray& axis,
                      bool keep_dim,
diff --git a/paddle/phi/ops/yaml/ops.yaml b/paddle/phi/ops/yaml/ops.yaml
@@ -1325,7 +1325,7 @@
   args : (Tensor x)
   output : Tensor
   infer_meta :
-    func : UnchangedInferMeta
+    func : DetInferMeta
   kernel :
     func : determinant
   backward : det_grad
diff --git a/python/paddle/tensor/creation.py b/python/paddle/tensor/creation.py
@@ -1360,7 +1360,9 @@ def eye(
 
     def _check_attr(attr, message):
         if isinstance(attr, ((Variable, core.eager.Tensor, paddle.pir.Value))):
-            assert len(attr.shape) == 1 and attr.shape[0] in [1, -1]
+            assert len(attr.shape) == 0 or (
+                len(attr.shape) == 1 and attr.shape[0] in [1, -1]
+            )
         elif not isinstance(attr, int) or attr < 0:
             raise TypeError(f"{message} should be a non-negative int.")
 
diff --git a/test/legacy_test/test_determinant_op.py b/test/legacy_test/test_determinant_op.py
@@ -91,13 +91,14 @@ def test_api_static(self):
         paddle.enable_static()
         with paddle.static.program_guard(paddle.static.Program()):
             x = paddle.static.data('X', self.shape)
-            out = paddle.linalg.det(x)
+            out_value = paddle.linalg.det(x)
             exe = paddle.static.Executor(self.place)
-            res = exe.run(feed={'X': self.x}, fetch_list=[out])
+            (out_np,) = exe.run(feed={'X': self.x}, fetch_list=[out_value])
         out_ref = np.linalg.det(self.x)
 
-        for out in res:
-            np.testing.assert_allclose(out, out_ref, rtol=0.001)
+        np.testing.assert_allclose(out_np, out_ref, rtol=0.001)
+        self.assertEqual(out_np.shape, out_ref.shape)
+        self.assertEqual(tuple(out_value.shape), out_ref.shape)
 
     def test_api_dygraph(self):
         paddle.disable_static(self.place)
diff --git a/test/legacy_test/test_zero_dim_no_backward_api.py b/test/legacy_test/test_zero_dim_no_backward_api.py
@@ -269,6 +269,16 @@ def test_matrix_rank(self):
         out_d = paddle.linalg.matrix_rank(d, tol=tol_2)
         self.assertEqual(out_d.shape, [2])
 
+    def test_eye_zero_dim_input(self):
+        # use zero-dim tensor as inputs
+        num_rows = paddle.to_tensor(5, stop_gradient=False)
+        num_cols = paddle.to_tensor(4, stop_gradient=False)
+        out = paddle.eye(num_rows, num_cols)
+
+        self.assertEqual(num_cols.shape, [])
+        self.assertEqual(num_rows.shape, [])
+        self.assertEqual(out.shape, [5, 4])
+
 
 class TestNoBackwardAPIStatic(unittest.TestCase):
     def setUp(self):
