[Zero-Dim] support ReshapeTransform/nll_loss/matmul support 0D (Paddl…

…ePaddle#53828)
zhwesky2010 · May 25, 2023 · 334bffd · 334bffd
1 parent 8f455ca
commit 334bffd
Show file tree

Hide file tree

Showing 22 changed files with 309 additions and 238 deletions.
diff --git a/paddle/fluid/operators/matmul_op.cc b/paddle/fluid/operators/matmul_op.cc
@@ -666,10 +666,6 @@ class MatMulOp : public framework::OperatorWithKernel {
       dim_out.resize(dim_out.size() - 1);
     }
 
-    if (dim_out.empty()) {
-      dim_out = {1};
-    }
-
     phi::DDim ddim_out = phi::make_ddim(dim_out);
 
     context->SetOutputDim("Out", ddim_out);

diff --git a/paddle/fluid/operators/matmul_v2_op.cc b/paddle/fluid/operators/matmul_v2_op.cc
@@ -91,9 +91,6 @@ void MatMulV2Op::InferShape(framework::InferShapeContext* ctx) const {
   if (!y_broadcasted) {
     new_dims.push_back(N);
   }
-  if (x_broadcasted && y_broadcasted) {
-    new_dims.push_back(1);
-  }
 
   ctx->SetOutputDim("Out", phi::make_ddim(new_dims));
   ctx->ShareLoD("X", "Out");

diff --git a/paddle/fluid/operators/mkldnn/matmul_mkldnn_op.cc b/paddle/fluid/operators/mkldnn/matmul_mkldnn_op.cc
@@ -334,8 +334,9 @@ void ExecuteMatMulV1(const ExecutionContext &ctx,
   matmul_p->execute(astream, matmul_args);
   astream.wait();
 
-  out->set_mem_desc(
-      dst_memory_p->get_desc().reshape(vectorize<int64_t>(out->dims())));
+  auto reshape_dims = out->dims().size() != 0 ? vectorize(out->dims())
+                                              : std::vector<int64_t>{1};
+  out->set_mem_desc(dst_memory_p->get_desc().reshape(reshape_dims));
 }
 
 template <typename T>

diff --git a/paddle/phi/backends/onednn/matmul_utils.h b/paddle/phi/backends/onednn/matmul_utils.h
@@ -146,8 +146,9 @@ inline void ExecuteMul(const OneDNNContext& dev_ctx,
   // This kernel is flattening dims so then we need to unflattened version
   // that should be set in out reshape require plain layout, but
   // MatmulV2MKLDNNHanlder enforces one so it should work
-  out->set_mem_desc(
-      dst_memory_p->get_desc().reshape(vectorize<int64_t>(out->dims())));
+  auto reshape_dims = out->dims().size() != 0 ? vectorize(out->dims())
+                                              : std::vector<int64_t>{1};
+  out->set_mem_desc(dst_memory_p->get_desc().reshape(reshape_dims));
 }
 
 template <typename T, typename T_out>
@@ -177,8 +178,9 @@ inline void ExecuteMatmul(const OneDNNContext& dev_ctx,
   matmul_p->execute(astream, matmul_args);
   astream.wait();
 
-  out->set_mem_desc(
-      dst_memory_p->get_desc().reshape(vectorize<int64_t>(out->dims())));
+  auto reshape_dims = out->dims().size() != 0 ? vectorize(out->dims())
+                                              : std::vector<int64_t>{1};
+  out->set_mem_desc(dst_memory_p->get_desc().reshape(reshape_dims));
 }
 
 }  // namespace funcs

diff --git a/paddle/phi/infermeta/backward.cc b/paddle/phi/infermeta/backward.cc
@@ -822,22 +822,22 @@ void NllLossGradInferMeta(const MetaTensor& x,
   if (check) {
     auto batch_size = x_dims[0];
     if (x_dims.size() == 2) {
-      PADDLE_ENFORCE_EQ(dout_dims.size(),
-                        1,
-                        phi::errors::InvalidArgument(
-                            "The dimensions of Input(Out@Grad) must be 1"));
       if (reduction == "none") {
+        PADDLE_ENFORCE_EQ(dout_dims.size(),
+                          1,
+                          phi::errors::InvalidArgument(
+                              "The dimensions of Input(Out@Grad) must be 1"));
         PADDLE_ENFORCE_EQ(
             dout_dims[0],
             batch_size,
             phi::errors::InvalidArgument(
                 "The unreduced size ofInput(Out@Grad) must be the "
                 "same as batch_size."));
       } else {
-        PADDLE_ENFORCE_EQ(dout_dims[0],
-                          1,
+        PADDLE_ENFORCE_EQ(dout_dims.size(),
+                          0,
                           phi::errors::InvalidArgument(
-                              "The reduced size of Input(Out@Grad) must be 1"));
+                              "The dimensions of Input(Out@Grad) must be 0"));
       }
     } else if (x_dims.size() == 4) {
       if (reduction == "none") {
@@ -855,10 +855,10 @@ void NllLossGradInferMeta(const MetaTensor& x,
                               "The dimensions of Input(Out@Grad) must be match "
                               "to Input(Label) dimensions."));
       } else {
-        PADDLE_ENFORCE_EQ(dout_dims[0],
-                          1,
+        PADDLE_ENFORCE_EQ(dout_dims.size(),
+                          0,
                           phi::errors::InvalidArgument(
-                              "The reduced size of Input(Out@Grad) must be 1"));
+                              "The dimensions of Input(Out@Grad) must be 0"));
       }
     }
   }

diff --git a/paddle/phi/infermeta/binary.cc b/paddle/phi/infermeta/binary.cc
@@ -2056,9 +2056,6 @@ void MatmulInferMeta(const MetaTensor& x,
   if (!y_broadcasted) {
     new_dims.push_back(N);
   }
-  if (x_broadcasted && y_broadcasted) {
-    new_dims.push_back(1);
-  }
 
   auto ddim_out = phi::make_ddim(new_dims);
 

diff --git a/paddle/phi/infermeta/ternary.cc b/paddle/phi/infermeta/ternary.cc
@@ -818,7 +818,7 @@ void NllLossRawInferMeta(const MetaTensor& input,
     if (reduction == "none") {
       out->set_dims({x_dims[0]});
     } else {
-      out->set_dims({1});
+      out->set_dims(phi::make_ddim({}));
     }
   } else if (x_dims.size() == 4) {
     PADDLE_ENFORCE_EQ(label_dims.size(),
@@ -841,10 +841,10 @@ void NllLossRawInferMeta(const MetaTensor& input,
     if (reduction == "none") {
       out->set_dims({x_dims[0], x_dims[2], x_dims[3]});
     } else {
-      out->set_dims({1});
+      out->set_dims(phi::make_ddim({}));
     }
   }
-  total_weight->set_dims({1});
+  total_weight->set_dims(phi::make_ddim({}));
   out->set_dtype(input.dtype());
   total_weight->set_dtype(input.dtype());
 }

diff --git a/paddle/phi/kernels/impl/matmul_kernel_impl.h b/paddle/phi/kernels/impl/matmul_kernel_impl.h
@@ -126,7 +126,7 @@ void MatMulFunctionImplWithBlas(
             M,
             N));
     VLOG(3) << "MatMul's case 1";
-    Out->Resize({1});
+    Out->Resize(phi::make_ddim({}));
     dev_ctx.template Alloc<T>(Out);
     blas.GEMM(CblasNoTrans,
               CblasTrans,
@@ -516,7 +516,7 @@ void MatMulFunctionImplWithCublasLt(
             N));
 
     // MatMul's case 0  =>  vector * vector
-    Out->Resize({1});
+    Out->Resize(phi::make_ddim({}));
     dev_ctx.template Alloc<T>(Out);
     VLOG(3) << "MatMul with blaslt case 1";
     blaslt::Run(dev_ctx,

diff --git a/paddle/phi/kernels/nll_loss_kernel.cc b/paddle/phi/kernels/nll_loss_kernel.cc
@@ -24,8 +24,7 @@ void NllLossKernel(const Context& dev_ctx,
                    const std::string& reduction,
                    DenseTensor* out) {
   DenseTensor total_weight;
-  total_weight.set_meta(
-      DenseTensorMeta(phi::CppTypeToDataType<T>::Type(), {1}));
+  total_weight.set_meta(DenseTensorMeta(phi::CppTypeToDataType<T>::Type(), {}));
   dev_ctx.template Alloc<T>(total_weight);
   NllLossRawKernel(dev_ctx,
                    input,

diff --git a/python/paddle/distribution/transform.py b/python/paddle/distribution/transform.py
@@ -856,8 +856,8 @@ class ReshapeTransform(Transform):
             #        [[[1., 1., 1.],
             #          [1., 1., 1.]]])
             print(reshape_transform.forward_log_det_jacobian(x))
-            # Tensor(shape=[1], dtype=float32, place=Place(gpu:0), stop_gradient=True,
-            #        [0.])
+            # Tensor(shape=[], dtype=float32, place=Place(gpu:0), stop_gradient=True,
+            #        0.)
     """
     _type = Type.BIJECTION
 
@@ -945,8 +945,7 @@ def _inverse_shape(self, shape):
         )
 
     def _forward_log_det_jacobian(self, x):
-        # TODO(zhouwei): should not set shape to [1], which is []
-        shape = x.shape[: x.dim() - len(self._in_event_shape)] or [1]
+        shape = x.shape[: x.dim() - len(self._in_event_shape)]
         return paddle.zeros(shape, dtype=x.dtype)
 
 

diff --git a/python/paddle/fluid/tests/unittests/distribution/test_distribution_transform.py b/python/paddle/fluid/tests/unittests/distribution/test_distribution_transform.py
@@ -1029,8 +1029,8 @@ def test_zerodim(self, input, expected):
 
         self.assertEqual(out.shape, [1, 1])
         self.assertEqual(reshape.inverse(out).shape, [])
-        # self.assertEqual(reshape.forward_log_det_jacobian(x).shape, [])
-        # self.assertEqual(reshape.inverse_log_det_jacobian(out).shape, [])
+        self.assertEqual(reshape.forward_log_det_jacobian(x).shape, [])
+        self.assertEqual(reshape.inverse_log_det_jacobian(out).shape, [])
         self.assertEqual(reshape.forward_shape(x.shape), (1, 1))
         self.assertEqual(reshape.inverse_shape(out.shape), ())
 

diff --git a/python/paddle/fluid/tests/unittests/test_matmul_op.py b/python/paddle/fluid/tests/unittests/test_matmul_op.py
@@ -77,12 +77,6 @@ def reference_matmul(X, Y, transpose_X=False, transpose_Y=False):
             Y = np.transpose(Y, tuple(dim))
 
     Out = np.matmul(X, Y)
-    if not Out.shape:
-        # We do not support 0-dimensional Tensors (scalars). So where
-        # np.matmul outputs a scalar, we must convert to a Tensor of
-        # shape (1, ) instead.
-        # Everywhere else, we are compatible with np.matmul.
-        Out = np.array([Out], dtype="float32")
     return Out
 
 
@@ -167,19 +161,14 @@ def test_out(self):
             with fluid.program_guard(fluid.Program()):
                 x = paddle.static.data(name="x", shape=[2], dtype="float64")
                 y = paddle.static.data(name='y', shape=[2], dtype='float64')
-                res = paddle.static.data(
-                    name="output", shape=[1], dtype="float64"
-                )
                 result = paddle.mm(x, y)
                 exe = fluid.Executor(fluid.CPUPlace())
                 data1 = np.random.rand(2)
                 data2 = np.random.rand(2)
                 np_res = exe.run(
                     feed={'x': data1, 'y': data2}, fetch_list=[result]
                 )
-                expected_result = np.matmul(
-                    data1.reshape(1, 2), data2.reshape(2, 1)
-                )
+                expected_result = np.matmul(data1, data2)
 
             np.testing.assert_allclose(
                 np_res,

diff --git a/python/paddle/fluid/tests/unittests/test_matmul_op_with_head.py b/python/paddle/fluid/tests/unittests/test_matmul_op_with_head.py
@@ -102,12 +102,6 @@ def reference_matmul_mul_head(
         Y = transpose_mat(Y)
 
     Out = matmul_head(X, Y, head_number)
-    if not Out.shape:
-        # We do not support 0-dimensional Tensors (scalars). So where
-        # np.matmul outputs a scalar, we must convert to a Tensor of
-        # shape (1, ) instead.
-        # Everywhere else, we are compatible with np.matmul.
-        Out = np.array([Out], dtype="float32")
     return Out
 
 
@@ -196,12 +190,6 @@ def reference_matmul_mul_head2(
         Y = transpose_mat(Y)
 
     Out = matmul_head2(X, Y, head_number)
-    if not Out.shape:
-        # We do not support 0-dimensional Tensors (scalars). So where
-        # np.matmul outputs a scalar, we must convert to a Tensor of
-        # shape (1, ) instead.
-        # Everywhere else, we are compatible with np.matmul.
-        Out = np.array([Out], dtype="float32")
     return Out
 
 

diff --git a/python/paddle/fluid/tests/unittests/test_matmul_v2_op.py b/python/paddle/fluid/tests/unittests/test_matmul_v2_op.py
@@ -45,12 +45,6 @@ def reference_matmul(X, Y, transpose_X=False, transpose_Y=False):
             Y = np.transpose(Y, tuple(dim))
 
     Out = np.matmul(X, Y)
-    if not Out.shape:
-        # We do not support 0-dimensional Tensors (scalars). So where
-        # np.matmul outputs a scalar, we must convert to a Tensor of
-        # shape (1, ) instead.
-        # Everywhere else, we are compatible with np.matmul.
-        Out = np.array([Out], dtype="float64")
     return Out