diff --git a/paddle/phi/kernels/impl/multi_dot_kernel_impl.h b/paddle/phi/kernels/impl/multi_dot_kernel_impl.h index ab0a08fe4b63d..3300f843c2d38 100644 --- a/paddle/phi/kernels/impl/multi_dot_kernel_impl.h +++ b/paddle/phi/kernels/impl/multi_dot_kernel_impl.h @@ -12,25 +12,11 @@ // See the License for the specific language governing permissions and // limitations under the License. -/* Copyright (c) 2021 PaddlePaddle Authors. All Rights Reserved. - -Licensed under the Apache License, Version 2.0 (the "License"); -you may not use this file except in compliance with the License. -You may obtain a copy of the License at - - http://www.apache.org/licenses/LICENSE-2.0 - -Unless required by applicable law or agreed to in writing, software -distributed under the License is distributed on an "AS IS" BASIS, -WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -See the License for the specific language governing permissions and -limitations under the License. */ - #pragma once #include "paddle/phi/core/dense_tensor.h" +#include "paddle/phi/kernels/full_kernel.h" #include "paddle/phi/kernels/funcs/blas/blas.h" - namespace phi { template @@ -195,6 +181,19 @@ void MultiDotKernel(const Context& dev_ctx, std::vector ins_dims(n); GetDims(ins, &ins_dims); + // If any numel is 0, then return. + bool size_0 = false; + for (size_t i = 0; i < n; i++) { + if (x[i]->numel() == 0) size_0 = true; + } + if (size_0) { + // For example: [2, 0], [0, 4] -> [2, 4] + if (out && out->numel() > 0) { + phi::Full( + dev_ctx, phi::IntArray(common::vectorize(out->dims())), 0, out); + } + return; + } const T scale = static_cast(1.0); if (n == 2) { auto mat_dim_a = phi::funcs::CreateMatrixDescriptor(ins_dims[0], 0, false); @@ -348,9 +347,23 @@ void MultiDotGradKernel(const Context& dev_ctx, auto blas = phi::funcs::GetBlas(dev_ctx); + bool size_0 = false; const auto n = ins.size(); for (size_t i = 0; i < n; i++) { dev_ctx.template Alloc(dx[i]); + + if (dx[i]->numel() == 0) { + size_0 = true; + } + } + if (size_0) { + for (size_t i = 0; i < n; i++) { + if (dx[i]->numel() > 0) { + phi::Full( + dev_ctx, phi::IntArray(common::vectorize(dx[i]->dims())), 0, dx[i]); + } + } + return; } std::vector ins_dims(n); diff --git a/test/legacy_test/test_multi_dot_op.py b/test/legacy_test/test_multi_dot_op.py index 308b9c1cc0652..2335ed31fb33e 100644 --- a/test/legacy_test/test_multi_dot_op.py +++ b/test/legacy_test/test_multi_dot_op.py @@ -31,14 +31,19 @@ def setUp(self): self.op_type = "multi_dot" self.python_api = paddle.linalg.multi_dot self.dtype = self.get_dtype() + self.init_shape() self.get_inputs_and_outputs() + def init_shape(self): + self.A_shape = (2, 8) + self.B_shape = (8, 4) + def get_dtype(self): return "float64" def get_inputs_and_outputs(self): - self.A = np.random.random((2, 8)).astype(self.dtype) - self.B = np.random.random((8, 4)).astype(self.dtype) + self.A = np.random.random(self.A_shape).astype(self.dtype) + self.B = np.random.random(self.B_shape).astype(self.dtype) self.inputs = {'X': [('x0', self.A), ('x1', self.B)]} self.outputs = {'Out': multi_dot([self.A, self.B])} @@ -55,6 +60,36 @@ def get_dtype(self): return "float16" +class TestMultiDotOp_ZeroSize1(TestMultiDotOp): + def get_inputs_and_outputs(self): + # result shape: [2, 3] + self.A = np.random.random((2, 10)).astype(self.dtype) + self.B = np.random.random((10, 0)).astype(self.dtype) + self.C = np.random.random((0, 3)).astype(self.dtype) + self.inputs = {'X': [('x0', self.A), ('x1', self.B), ('x2', self.C)]} + self.outputs = {'Out': multi_dot([self.A, self.B, self.C])} + + def test_check_grad(self): + self.check_grad(['x0'], 'Out', check_pir=True) + self.check_grad(['x1'], 'Out', check_pir=True) + self.check_grad(['x2'], 'Out', check_pir=True) + + +class TestMultiDotOp_ZeroSize2(TestMultiDotOp): + def get_inputs_and_outputs(self): + # result shape: [0, 3] + self.A = np.random.random((0, 10)).astype(self.dtype) + self.B = np.random.random((10, 4)).astype(self.dtype) + self.C = np.random.random((4, 3)).astype(self.dtype) + self.inputs = {'X': [('x0', self.A), ('x1', self.B), ('x2', self.C)]} + self.outputs = {'Out': multi_dot([self.A, self.B, self.C])} + + def test_check_grad(self): + self.check_grad(['x0'], 'Out', check_pir=True) + self.check_grad(['x1'], 'Out', check_pir=True) + self.check_grad(['x2'], 'Out', check_pir=True) + + @unittest.skipIf( not core.is_compiled_with_cuda() or not core.is_bfloat16_supported(core.CUDAPlace(0)),