fix add infermata (#66875)

wanghuancoder · web-flow · commit bcb987995a23 · 2024-08-01T19:07:20.000+08:00
* fix add infermata

* refine

* refine
diff --git a/paddle/phi/infermeta/binary.cc b/paddle/phi/infermeta/binary.cc
@@ -1647,8 +1647,9 @@ void DotInferMeta(const MetaTensor& x, const MetaTensor& y, MetaTensor* out) {
 
 void ElementwiseInferMeta(const MetaTensor& x,
                           const MetaTensor& y,
-                          MetaTensor* out) {
-  return ElementwiseRawInferMeta(x, y, -1, out);
+                          MetaTensor* out,
+                          MetaConfig config) {
+  return ElementwiseRawInferMeta(x, y, -1, out, config);
 }
 
 void BitwiseShiftInferMeta(const MetaTensor& x,
@@ -1691,6 +1692,7 @@ void ElementwiseRawInferMeta(const MetaTensor& x,
     std::vector<int> x_dims_array(max_dim);
     std::vector<int> y_dims_array(max_dim);
     std::vector<int> out_dims_array(max_dim);
+
 #ifdef PADDLE_WITH_DNNL
     bool should_rotate =
         config.is_run_mkldnn_kernel &&
@@ -1699,7 +1701,7 @@ void ElementwiseRawInferMeta(const MetaTensor& x,
         (x_dims.size() >= 3 || y_dims.size() >= 3);
     if (should_rotate) {
       // Pick bigger shape and rotate this one
-      bool x_over_y = (x_dims.size() > y_dims.size());
+      bool x_over_y = (common::product(x_dims) > common::product(y_dims));
       auto vdims = x_over_y ? common::vectorize<int>(x_dims)
                             : common::vectorize<int>(y_dims);
       std::rotate(vdims.begin() + 1, vdims.begin() + 2, vdims.end());
diff --git a/paddle/phi/infermeta/binary.h b/paddle/phi/infermeta/binary.h
@@ -324,7 +324,8 @@ void DropoutNdInferMeta(const MetaTensor& x,
 
 TEST_API void ElementwiseInferMeta(const MetaTensor& x,
                                    const MetaTensor& y,
-                                   MetaTensor* out);
+                                   MetaTensor* out,
+                                   MetaConfig config = MetaConfig());
 
 void ElementwiseRawInferMeta(const MetaTensor& x_meta,
                              const MetaTensor& y_meta,
