fix 0 size Tensor for expand kernel in onednn

Author: DanielSun11

paddle/phi/kernels/onednn/expand_kernel.cc

@@ -39,16 +39,51 @@ void ExpandKernel(const Context& dev_ctx,
   auto x_vec_dims = common::vectorize(x.dims());
 
   auto out_new_dims = shape.GetData();
+  bool has_zero_size = false;
 
   for (size_t i = 0; i < out_new_dims.size(); ++i) {
-    out_new_dims[i] = out_new_dims[i] > 0 ? out_new_dims[i] : x_vec_dims[i];
+    out_new_dims[i] = out_new_dims[i] >= 0 ? out_new_dims[i] : x_vec_dims[i];
   }
 
   if (x_vec_dims.size() != out_new_dims.size()) {
     x_vec_dims = GetExtendedXDims(x_vec_dims, out_new_dims.size());  // NOLINT
   }
 
+  for (size_t i = 0; i < x_vec_dims.size(); ++i) {
+    PADDLE_ENFORCE_GE(
+        out_new_dims[i],
+        0,
+        common::errors::InvalidArgument(
+            "The expanded size (%d) for non-existing dimensions must be "
+            "positive for expand_v2 op.",
+            out_new_dims[i]));
+            
+    PADDLE_ENFORCE_GE(
+        x_vec_dims[i],
+        0,
+        common::errors::InvalidArgument(
+            "The expanded size (%d) for non-existing dimensions must be "
+            "positive for expand_v2 op.",
+            x_vec_dims[i]));
+
+    PADDLE_ENFORCE_EQ(
+        x_vec_dims[i] == 1 || x_vec_dims[i] == out_new_dims[i],
+        true,
+        common::errors::InvalidArgument(
+            "The value (%d) of the non-singleton dimension does not match"
+            " the corresponding value (%d) in shape for expand_v2 op.",
+            x_vec_dims[i],
+            out_new_dims[i]));
+    if (out_new_dims[i] == 0) {
+      has_zero_size = true;
+    }
+  }
+
   out->Resize(common::make_ddim(out_new_dims));
+  if (has_zero_size) {
+    dev_ctx.template Alloc<T>(out);
+    return;
+  }
   funcs::BroadcastDataOneDNNHandler<T> handler(dnnl::algorithm::binary_add,
                                                onednn_engine,
                                                dev_ctx.GetPlace(),

test/legacy_test/test_expand_v2_op.py

@@ -684,6 +684,48 @@ def test_value_list_shape2(self):
                 x = paddle.expand(x, shape=[shape1, 1, -1, -1])
                 np.testing.assert_equal(tuple(x.shape), (-1, 1, -1, -1))
 
+class TestExpandV2OneDNNOp(OpTest):
+    def setUp(self):
+        self.op_type = "expand_v2"
+        self.init_data()
+        self.x = np.random.random(self.ori_shape).astype("float32")
+        self.attrs = {'shape': self.shape, 'use_mkldnn': True}
+        self.set_inputs()
+        self.set_additional_inputs()
+        output = np.zeros(self.expect_shape).astype("float32")
+        self.outputs = {'Out': output}
+
+    def set_inputs(self):
+        self.inputs = {'X': self.x}
+
+    def set_additional_inputs(self):
+        pass
+
+    def init_data(self):
+        self.ori_shape = [1, 1, 1, 140]
+        self.shape = [2, 3, 0, 140]
+        self.expect_shape = [2, 3, 0, 140]
+
+    def test_check_output(self):
+        self.check_output_with_place(core.CPUPlace(), check_pir_onednn=True,check_dygraph=False)
+        
+    # def test_check_grad(self):
+    #     self.check_grad_with_place(
+    #         core.CPUPlace(), ["X"], "Out", check_pir_onednn=True, check_dygraph=False
+    #     )
+class TestExpandV2ZeroSizeOneDNNOp(TestExpandV2OneDNNOp):
+
+    def init_data(self):
+        self.ori_shape = (1, 3)
+        self.shape = (0, 3)
+        self.expect_shape = (0, 3)
+
+class TestExpandV2ZeroSizeOneDNNOp2(TestExpandV2OneDNNOp):
+
+    def init_data(self):
+        self.ori_shape = (1, 3)
+        self.shape = (1, 0, 3)
+        self.expect_shape = (1, 0, 3)
 
 if __name__ == "__main__":
     paddle.enable_static()