Add broadcast_shape api (#28257)

* add broadcast_shape api

* add ut

* follow comments

* add example code, test=dodument_fix

* update example code, test=document_fix

paddle/fluid/operators/common_infer_shape_functions.cc

@@ -13,6 +13,7 @@ See the License for the specific language governing permissions and
 limitations under the License. */
 
 #include "paddle/fluid/operators/common_infer_shape_functions.h"
+
 #include <algorithm>
 #include <vector>
 
@@ -28,6 +29,7 @@ class InferShapeContext;
 namespace paddle {
 namespace operators {
 namespace details {
+
 inline void GetBroadcastDimsArrays(const framework::DDim &x_dims,
                                    const framework::DDim &y_dims,
                                    int *x_dims_array, int *y_dims_array,
@@ -76,6 +78,20 @@ inline void GetBroadcastDimsArrays(const framework::DDim &x_dims,
     }
   }
 }
+
+framework::DDim BroadcastTwoDims(const framework::DDim &x_dims,
+                                 const framework::DDim &y_dims, int axis) {
+  int max_dim = std::max(x_dims.size(), y_dims.size());
+  axis = (axis == -1 ? std::abs(x_dims.size() - y_dims.size()) : axis);
+  std::vector<int> x_dims_array(max_dim);
+  std::vector<int> y_dims_array(max_dim);
+  std::vector<int> out_dims_array(max_dim);
+  GetBroadcastDimsArrays(x_dims, y_dims, x_dims_array.data(),
+                         y_dims_array.data(), out_dims_array.data(), max_dim,
+                         axis);
+  return framework::make_ddim(out_dims_array);
+}
+
 }  // namespace details
 
 // shape input(0) -> output(0) without change.
@@ -153,16 +169,9 @@ void BinaryOpBroadcastInferShape(framework::InferShapeContext *ctx) {
     ctx->ShareDim(x_name, /*->*/ out_name);
     ctx->ShareLoD(x_name, /*->*/ out_name);
   } else {
-    int max_dim = std::max(x_dims.size(), y_dims.size());
     int axis = ctx->Attrs().Get<int>("axis");
-    axis = (axis == -1 ? std::abs(x_dims.size() - y_dims.size()) : axis);
+    auto out_dims = details::BroadcastTwoDims(x_dims, y_dims, axis);
-    std::vector<int> x_dims_array(max_dim);
+    ctx->SetOutputDim(out_name, out_dims);
-    std::vector<int> y_dims_array(max_dim);
-    std::vector<int> out_dims_array(max_dim);
-    details::GetBroadcastDimsArrays(x_dims, y_dims, x_dims_array.data(),
-                                    y_dims_array.data(), out_dims_array.data(),
-                                    max_dim, axis);
-    ctx->SetOutputDim(out_name, framework::make_ddim(out_dims_array));
     ctx->ShareLoD(x_name, /*->*/ out_name);
   }
 }

paddle/fluid/operators/common_infer_shape_functions.h

@@ -28,7 +28,10 @@ class InferShapeContext;
 
 namespace paddle {
 namespace operators {
-
+namespace details {
+framework::DDim BroadcastTwoDims(const framework::DDim& x_dims,
+                                 const framework::DDim& y_dims, int axis = -1);
+}
 // shape input(0) -> output(0) without change.
 void UnaryOpUnchangedInferShape(framework::InferShapeContext* ctx);
 // shape input(0) -> output(0) without change, check if axis in range [-Rank(x),

paddle/fluid/pybind/pybind.cc

@@ -54,6 +54,7 @@ limitations under the License. */
 #include "paddle/fluid/memory/allocation/allocator_strategy.h"
 #include "paddle/fluid/memory/allocation/mmap_allocator.h"
 #include "paddle/fluid/operators/activation_op.h"
+#include "paddle/fluid/operators/common_infer_shape_functions.h"
 #include "paddle/fluid/operators/py_func_op.h"
 #include "paddle/fluid/platform/cpu_helper.h"
 #include "paddle/fluid/platform/cpu_info.h"
@@ -467,6 +468,12 @@ PYBIND11_MODULE(core_noavx, m) {
             << ", sci_mode=" << print_opt.sci_mode;
   });
 
+  m.def("broadcast_shape", [](const std::vector<int64_t> &x_dim,
+                              const std::vector<int64_t> &y_dim) {
+    return vectorize(operators::details::BroadcastTwoDims(
+        make_ddim(x_dim), make_ddim(y_dim), -1));
+  });
+
   m.def(
       "_append_python_callable_object_and_return_id",
       [](py::object py_obj) -> size_t {

python/paddle/__init__.py

@@ -200,6 +200,8 @@ from .tensor.math import isfinite  #DEFINE_ALIAS
 from .tensor.math import isinf  #DEFINE_ALIAS
 from .tensor.math import isnan  #DEFINE_ALIAS
 from .tensor.math import prod  #DEFINE_ALIAS
+from .tensor.math import broadcast_shape  #DEFINE_ALIAS
+
 from .tensor.random import multinomial  #DEFINE_ALIAS
 from .tensor.random import standard_normal
 from .tensor.random import normal
@@ -220,7 +222,7 @@ from .tensor.search import index_select  #DEFINE_ALIAS
 from .tensor.search import nonzero  #DEFINE_ALIAS
 from .tensor.search import sort  #DEFINE_ALIAS
 
-from .tensor.to_string import set_printoptions
+from .tensor.to_string import set_printoptions  #DEFINE_ALIAS
 
 from .framework.random import seed  #DEFINE_ALIAS
 from .framework.random import get_cuda_rng_state  #DEFINE_ALIAS

python/paddle/fluid/tests/unittests/test_broadcast_shape.py

@@ -0,0 +1,35 @@
+# Copyright (c) 2020 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.
+
+import unittest
+import numpy as np
+import paddle
+
+
+class TestBroadcastShape(unittest.TestCase):
+    def test_result(self):
+        shape = paddle.broadcast_shape([2, 1, 3], [1, 3, 1])
+        self.assertEqual(shape, [2, 3, 3])
+
+        shape = paddle.broadcast_shape(
+            [-1, 1, 3], [1, 3, 1])  #support compile time infershape
+        self.assertEqual(shape, [-1, 3, 3])
+
+    def test_error(self):
+        self.assertRaises(ValueError, paddle.broadcast_shape, [2, 1, 3],
+                          [3, 3, 1])
+
+
+if __name__ == "__main__":
+    unittest.main()

python/paddle/tensor/__init__.py

@@ -164,6 +164,8 @@ from .math import isnan  #DEFINE_ALIAS
 from .math import prod  #DEFINE_ALIAS
 from .math import all  #DEFINE_ALIAS
 from .math import any  #DEFINE_ALIAS
+from .math import broadcast_shape  #DEFINE_ALIAS
+
 from .random import multinomial  #DEFINE_ALIAS
 from .random import standard_normal
 from .random import normal
@@ -194,4 +196,4 @@ from .stat import median  #DEFINE_ALIAS
 # from .tensor import Tensor        #DEFINE_ALIAS
 # from .tensor import LoDTensor        #DEFINE_ALIAS
 # from .tensor import LoDTensorArray        #DEFINE_ALIAS
-from .to_string import set_printoptions
+from .to_string import set_printoptions  #DEFINE_ALIAS

python/paddle/tensor/math.py

@@ -121,7 +121,8 @@ __all__ = [
         'kron',
         'isfinite',
         'isinf',
-        'isnan'
+        'isnan',
+        'broadcast_shape'
 ]
 # yapf: enable.
 
@@ -2133,3 +2134,30 @@ def any(x, axis=None, keepdim=False, name=None):
         outputs={'Out': out},
         attrs=attrs)
     return out
+
+def broadcast_shape(x_shape, y_shape):
+    """
+    The function returns the shape of doing operation with broadcasting on tensors of x_shape and y_shape, please refer to :ref:`user_guide_broadcasting` for more details.
+
+    Args:
+        x_shape (list[int]|tuple[int]): A shape of tensor.
+        y_shape (list[int]|tuple[int]): A shape of tensor.
+        
+
+    Returns:
+        list[int], the result shape.
+
+    Examples:
+        .. code-block:: python
+
+            import paddle
+
+            shape = paddle.broadcast_shape([2, 1, 3], [1, 3, 1])
+            # [2, 3, 3]
+            
+            # shape = paddle.broadcast_shape([2, 1, 3], [3, 3, 1])
+            # ValueError (terminated with error message).
+
+    """
+
+    return core.broadcast_shape(x_shape, y_shape)