diff --git a/paddle/phi/kernels/impl/matmul_kernel_impl.h b/paddle/phi/kernels/impl/matmul_kernel_impl.h index f7b86ea8bf4045..d087e25d1e1fa3 100644 --- a/paddle/phi/kernels/impl/matmul_kernel_impl.h +++ b/paddle/phi/kernels/impl/matmul_kernel_impl.h @@ -2007,16 +2007,36 @@ void MatmulKernel(const Context& ctx, bool transpose_x, bool transpose_y, DenseTensor* out) { - PADDLE_ENFORCE_NE(common::product(x.dims()), - 0, - common::errors::InvalidArgument( - "The Input(X) dims size must not be equal " - "0, but received dims size is 0.")); - PADDLE_ENFORCE_NE(common::product(y.dims()), - 0, - common::errors::InvalidArgument( - "The Input(Y) dims size must not be equal " - "0, but received dims size is 0.")); + if (x.numel() == 0 || y.numel() == 0) { + auto x_dims = x.dims(); + auto y_dims = y.dims(); + if (transpose_x) { + std::swap(x_dims[x_dims.size() - 1], x_dims[x_dims.size() - 2]); + } + if (transpose_y) { + std::swap(y_dims[y_dims.size() - 1], y_dims[y_dims.size() - 2]); + } + std::vector out_dims(x_dims.size() - 1 + y_dims.size() - 1); + for (int64_t i = 0; i < x_dims.size() - 1; ++i) { + out_dims[i] = x_dims[i]; + } + for (int64_t i = 1; i < y_dims.size(); ++i) { + out_dims[x_dims.size() - 1 + i - 1] = y_dims[i]; + } + out->Resize(phi::make_ddim(out_dims)); + ctx.template Alloc(out); + return; + } + PADDLE_ENFORCE_GE( + common::product(x.dims()), + 0, + common::errors::InvalidArgument( + "The dims of Input(X) should be greater than or equal to 0.")); + PADDLE_ENFORCE_GE( + common::product(y.dims()), + 0, + common::errors::InvalidArgument( + "The dims of Input(Y) should be greater than or equal to 0.")); const std::vector x_dims = common::vectorize(x.dims()); const std::vector y_dims = common::vectorize(y.dims()); MatmulJudgeDtypeKernel( diff --git a/test/legacy_test/test_tensordot.py b/test/legacy_test/test_tensordot.py index 3d15a0f7931b4d..1df90d2a226267 100644 --- a/test/legacy_test/test_tensordot.py +++ b/test/legacy_test/test_tensordot.py @@ -371,5 +371,116 @@ def set_dtype(self): self.dtype = np.float64 +class TestTensordotAPIZeroSize(TestTensordotAPI): + def set_input_shape(self): + self.x_shape = [0, 5, 5, 5] + self.y_shape = [0, 5, 5, 5] + + def set_input_data(self): + self.x = np.random.random(self.x_shape).astype(self.dtype) + self.y = np.random.random(self.y_shape).astype(self.dtype) + + def set_test_axes(self): + self.all_axes = [ + [[], []], + ] + + +class TestTensordotAPIFloat64ZeroSize(TestTensordotAPIZeroSize): + def set_dtype(self): + self.dtype = np.float64 + + +class TestTensordotAPIZeroSize(TestTensordotAPI): + def set_input_shape(self): + self.x_shape = [0, 5, 5, 5] + self.y_shape = [0, 5, 5, 5] + + def set_input_data(self): + self.x = np.random.random(self.x_shape).astype(self.dtype) + self.y = np.random.random(self.y_shape).astype(self.dtype) + + def set_test_axes(self): + self.all_axes = [ + [[], []], + ] + + def set_dtype(self): + self.dtype = np.float64 + + +class TestTensordotAPIZeroSizeMultipleDims1(TestTensordotAPI): + def set_input_shape(self): + self.x_shape = [0, 0, 5, 5] + self.y_shape = [0, 0, 5, 5] + + def set_test_axes(self): + self.all_axes = [ + [[], []], + ] + + +class TestTensordotAPIZeroSizeMultipleDims2(TestTensordotAPI): + def set_input_shape(self): + self.x_shape = [5, 0, 5, 0] + self.y_shape = [5, 0, 5, 0] + + def set_test_axes(self): + self.all_axes = [ + [[], []], + ] + + +class TestTensordotAPIZeroSizeDifferentDims1(TestTensordotAPI): + def set_input_shape(self): + self.x_shape = [5, 5, 0, 5] + self.y_shape = [5, 5, 0, 5] + + def set_test_axes(self): + self.all_axes = [ + [[], []], + ] + + +class TestTensordotAPIZeroSizeDifferentDims2(TestTensordotAPI): + def set_input_shape(self): + self.x_shape = [5, 5, 5, 0] + self.y_shape = [5, 5, 5, 0] + + def set_test_axes(self): + self.all_axes = [ + [[], []], + ] + + +class TestTensordotAPISingleElementAndZeroSize(TestTensordotAPI): + def set_input_shape(self): + self.x_shape = [1, 5, 5, 5] + self.y_shape = [0, 5, 5, 5] + + def set_test_axes(self): + self.all_axes = [ + [[], []], + ] + + +class TestBroadcastWithZeroSize1(unittest.TestCase): + def setUp(self): + self.x_shape = [5, 0, 3] + self.y_shape = [3, 4, 0] + + def set_test_axes(self): + self.all_axes = [[], []] + + +class TestBroadcastWithZeroSize2(unittest.TestCase): + def setUp(self): + self.x_shape = [5, 0, 3] + self.y_shape = [3, 0] + + def set_test_axes(self): + self.all_axes = [[], []] + + if __name__ == "__main__": unittest.main()