[cherrypick][inference Zero-Dim] Support 0-Dim Tensor in Paddle-TensorRT (#53752)

* scale, square, sum, swish trt op converter support zero dim (#53660)

* [Paddle-Inference] Support trt 0dims of expand_as_v2 and mish. (#53627)

* support_expand_mish

* add unitest for reshpe 0 dims (#53685)

* Add trt pow converter. (#53462)

* Add trt pow converter.

* update to use AddConstantLayer

* add dims=0 ut

* [inference Zero-Dim]add equal, elementwise_op trt 0d (#53704)

* [inference Zero-Dim]prelu trt converter support zero dim tensor (#53634)

* prelu op trt converter support zero dim

* [Inference Zero-Dim] Support trt 0dim of gelu, hard_swish, hard_sigmoid and leaky_relu (#53714)

* support_act
* delete_silu

* [inference zero dim] softmax, stack op trt converter support zero dim (#53729)

* softmax support

* support stack

* remove unused code

* update

---------
Co-authored-by: NYuanle Liu <yuanlehome@163.com>
Co-authored-by: Nxiaoxiaohehe001 <49090790+xiaoxiaohehe001@users.noreply.github.com>
Co-authored-by: Nzhoutianzi666 <39978853+zhoutianzi666@users.noreply.github.com>
Co-authored-by: NWilber <jiweibo@baidu.com>

paddle/fluid/framework/ir/trt_support_nhwc_pass.cc

@@ -356,6 +356,8 @@ void TrtSupportNHWCPass::ApplyImpl(Graph *graph) const {
     }
   };
   InsertTransposeOp();
+
+  AddStatis(transposed_ops.size());
 }
 
 }  // namespace ir

paddle/fluid/inference/tensorrt/convert/elementwise_op.cc

@@ -319,6 +319,37 @@ class ElementwiseTensorModOpConverter : public ElementwiseTensorOpConverter {
  public:
   ElementwiseTensorModOpConverter() { op_type_ = "mod"; }
 };
+
+// The diff between `pow` and `elementwise_pow` is in:
+// https://github.com/PaddlePaddle/Paddle/blob/release/2.4/python/paddle/tensor/math.py#L420
+class PowOpConverter : public OpConverter {
+ public:
+  PowOpConverter() {}
+  void operator()(const framework::proto::OpDesc& op,
+                  const framework::Scope& scope,
+                  bool test_mode) override {
+    VLOG(3) << "Convert a pow op to TensorRT IElementWiseLayer";
+    framework::OpDesc op_desc(op, nullptr);
+    auto* X = engine_->GetITensor(op_desc.Input("X").front());
+    float factor = PADDLE_GET_CONST(float, op_desc.GetAttr("factor"));
+    nvinfer1::Dims dims_x = X->getDimensions();
+    auto output_name = op_desc.Output("Out")[0];
+
+    nvinfer1::Dims trt_dims_y;
+    trt_dims_y.nbDims = dims_x.nbDims;
+    for (int i = 0; i < trt_dims_y.nbDims; i++) {
+      trt_dims_y.d[i] = 1;
+    }
+
+    std::vector<float> w_data{factor};
+    auto* Y = AddConstantLayer(w_data.data(), trt_dims_y);
+
+    auto* layer = TRT_ENGINE_ADD_LAYER(
+        engine_, ElementWise, *X, *Y, nvinfer1::ElementWiseOperation::kPOW);
+    RreplenishLayerAndOutput(layer, "elementwise", {output_name}, test_mode);
+  }
+};
+
 }  // namespace tensorrt
 }  // namespace inference
 }  // namespace paddle
@@ -369,3 +400,5 @@ REGISTER_TRT_OP_CONVERTER(logical_and, ElementwiseTensorLogicalAndOpConverter);
 REGISTER_TRT_OP_CONVERTER(less_equal, ElementwiseTensorLessEqualOpConverter);
 REGISTER_TRT_OP_CONVERTER(greater_equal,
                           ElementwiseTensorGreaterEqualOpConverter);
+
+REGISTER_TRT_OP_CONVERTER(pow, PowOpConverter);

paddle/fluid/inference/tensorrt/convert/expand_v2_op.cc

@@ -83,7 +83,10 @@ class ExpandOpConverter : public OpConverter {
       input_shape_tensor = Shape(input);
     }
 
-    auto* newInputTensor = Reshape(input, input_shape_tensor);
+    auto* newInputTensor =
+        Reshape(input,
+                input_shape_tensor,
+                ("expand_v2: reshape: (Output(" + output_name + ")").c_str());
 
     std::vector<int32_t> start_vec(shape_rank, 0);
     nvinfer1::Dims start;

paddle/fluid/inference/tensorrt/convert/op_converter.h

@@ -407,13 +407,11 @@ class OpConverter {
   }
 
   nvinfer1::ITensor* Reshape(nvinfer1::ITensor* input,
-                             nvinfer1::ITensor* newShape) {
-    nvinfer1::ITensor* oldShape = Shape(input);
-    if (oldShape == newShape) {
-      return input;
-    }
+                             nvinfer1::ITensor* newShape,
+                             const std::string& name = "reshape") {
     auto* shuffle = TRT_ENGINE_ADD_LAYER(engine_, Shuffle, *input);
     shuffle->setInput(1, *newShape);
+    shuffle->setName(name.c_str());
     return shuffle->getOutput(0);
   }
 

paddle/fluid/inference/tensorrt/convert/prelu_op.cc

@@ -87,7 +87,6 @@ class PReluOpConverter : public OpConverter {
           if (hw_tensor != nullptr) {
             shape_tensor = Concat(
                 std::vector<nvinfer1::ITensor*>{n_tensor, c_tensor, hw_tensor});
-
           } else {
             shape_tensor =
                 Concat(std::vector<nvinfer1::ITensor*>{n_tensor, c_tensor});

paddle/fluid/inference/tensorrt/convert/softmax_op.cc

@@ -38,8 +38,23 @@ class SoftMaxOpConverter : public OpConverter {
                    ? PADDLE_GET_CONST(int, op_desc.GetAttr("axis"))
                    : -1;
 
-    auto* layer = TRT_ENGINE_ADD_LAYER(
-        engine_, SoftMax, *const_cast<nvinfer1::ITensor*>(input1));
+    // support 0 or 1 dims input
+    bool is_0_dims = input_dims == 0;
+    bool is_1_dims = input_dims == 1;
+    if (is_0_dims || is_1_dims) {
+      auto reshaped_layer = TRT_ENGINE_ADD_LAYER(engine_, Shuffle, *input1);
+      nvinfer1::Dims reshaped_dims;
+      reshaped_dims.nbDims = 2;
+      reshaped_dims.d[0] = 1;
+      reshaped_dims.d[1] = is_0_dims ? 1 : input_shape.d[0];
+      reshaped_layer->setReshapeDimensions(reshaped_dims);
+      input1 = reshaped_layer->getOutput(0);
+      input_shape = input1->getDimensions();
+      input_dims = input_shape.nbDims;
+      axis = -1;
+    }
+
+    auto* layer = TRT_ENGINE_ADD_LAYER(engine_, SoftMax, *input1);
     uint32_t axes = std::max(0, input_dims - 3);
     // TODO(cryoco): Poor workaround. Fix padded dims problem when TRT layers
     // support Nd.
@@ -68,11 +83,22 @@ class SoftMaxOpConverter : public OpConverter {
       }
     }
     layer->setAxes(1 << axes);
-    auto output_name = op_desc.Output("Out")[0];
-    RreplenishLayerAndOutput(layer, "softmax", {output_name}, test_mode);
 
     // The trt will not run int for softmax.
     engine_->SetTensorDynamicRange(input1, 1.0);
+    auto output_name = op_desc.Output("Out")[0];
+
+    // support 0 or 1 dims input
+    if (is_0_dims || is_1_dims) {
+      auto reshaped_layer =
+          TRT_ENGINE_ADD_LAYER(engine_, Shuffle, *layer->getOutput(0));
+      reshaped_layer->setReshapeDimensions(
+          engine_->GetITensor(op_desc.Input("X")[0])->getDimensions());
+      RreplenishLayerAndOutput(
+          reshaped_layer, "reshape_softmax_reshape", {output_name}, test_mode);
+    } else {
+      RreplenishLayerAndOutput(layer, "softmax", {output_name}, test_mode);
+    }
   }
 };
 

paddle/fluid/inference/tensorrt/convert/stack_op.cc

@@ -65,12 +65,11 @@ class StackOpConverter : public OpConverter {
     auto* after_shape_tensor = Concat(shape_tensor_vec);
 
     for (int i = 0; i < input_num; ++i) {
-      auto* reshape_layer = TRT_ENGINE_ADD_LAYER(engine_, Shuffle, *inputs[i]);
-      reshape_layer->setInput(1, *after_shape_tensor);
-      inputs[i] = reshape_layer->getOutput(0);
-      reshape_layer->setName(("stack: reshape: (Output( " + std::to_string(i) +
-                              " )" + output_name + ")")
-                                 .c_str());
+      inputs[i] = Reshape(inputs[i],
+                          after_shape_tensor,
+                          ("stack: reshape: (Output( " + std::to_string(i) +
+                           " )" + output_name + ")")
+                              .c_str());
     }
 
     auto* layer = TRT_ENGINE_ADD_LAYER(

paddle/fluid/inference/tensorrt/op_teller.cc

@@ -105,15 +105,34 @@ struct SimpleOpTypeSetTeller : public Teller {
         "erf",        "floor",       "round",
         "sign",       "silu",        "logical_not",
         "reciprocal", "tanh_shrink", "logsigmoid",
-        "rsqrt"};
+        "rsqrt",      "swish",       "hard_sigmoid",
+        "hard_swish", "leaky_relu"};
     std::unordered_set<std::string> unary_list = {
-        "exp",        "log",  "sqrt",        "abs",        "sin",
-        "cos",        "tan",  "tanh",        "sinh",       "cosh",
-        "asin",       "acos", "atan",        "asinh",      "acosh",
-        "atanh",      "ceil", "celu",        "floor",      "round",
-        "sign",       "silu", "logical_not", "reciprocal", "tanh_shrink",
-        "logsigmoid", "erf",  "bitwise_not", "equal",      "not_equal",
-        "rsqrt"};
+        "exp",   "log",         "sqrt",       "abs",         "sin",
+        "cos",   "tan",         "tanh",       "sinh",        "cosh",
+        "asin",  "acos",        "atan",       "asinh",       "acosh",
+        "atanh", "ceil",        "celu",       "floor",       "round",
+        "sign",  "logical_not", "reciprocal", "tanh_shrink", "logsigmoid",
+        "erf",   "bitwise_not", "equal",      "not_equal",   "rsqrt"};
+
+    // Static shape does not support 0 or 1 dim's input.
+    if (!with_dynamic_shape) {
+      auto inputs = desc.Inputs();
+      for (auto iter : inputs) {
+        for (auto var_name : iter.second) {
+          auto* block = desc.Block();
+          if (block) {
+            auto* var_desc = block->FindVar(var_name);
+            // Can't get feed op's TensorDesc
+            if (op_type != "feed" && var_desc && !var_desc->Persistable()) {
+              const auto shape = var_desc->GetShape();
+              if (shape.size() == 1 || shape.size() == 0) return false;
+            }
+          }
+        }
+      }
+    }
+
     if (act_op_list.find(op_type) != act_op_list.end()) {
       auto* block = desc.Block();
       if (block == nullptr) {
@@ -122,15 +141,6 @@ struct SimpleOpTypeSetTeller : public Teller {
                    "the pass.";
         return false;
       }
-      auto x_var_name = desc.Input("X")[0];
-      auto* x_var_desc = block->FindVar(x_var_name);
-      const auto x_shape = x_var_desc->GetShape();
-      if (!with_dynamic_shape && (x_shape.size() == 1 || x_shape.size() == 0)) {
-        VLOG(3) << op_type
-                << " op does not support input's dim is 1 or 0 in tensorrt "
-                   "static shape mode.";
-        return false;
-      }
 #if !IS_TRT_VERSION_GE(7000)
       if (op_type == "erf") {
         VLOG(3) << op_type << " op does not support tensorrt.";
@@ -138,6 +148,9 @@ struct SimpleOpTypeSetTeller : public Teller {
       }
 #endif
 #if !IS_TRT_VERSION_GE(8600)
+      auto x_var_name = desc.Input("X")[0];
+      auto* x_var_desc = block->FindVar(x_var_name);
+      const auto x_shape = x_var_desc->GetShape();
       if (x_shape.size() == 0 && unary_list.find(op_type) != unary_list.end()) {
         VLOG(3) << op_type
                 << " op does not support 0 dim input when TensorRT < 8.6.";
@@ -145,24 +158,6 @@ struct SimpleOpTypeSetTeller : public Teller {
       }
 #endif
     }
-    // In static shape in Paddle-TRT, we can't allow that one op has a
-    // 1D intermediate tensor as input.
-    if (!with_dynamic_shape) {
-      auto inputs = desc.Inputs();
-      for (auto iter : inputs) {
-        for (auto var_name : iter.second) {
-          auto* block = desc.Block();
-          if (block) {
-            auto* var_desc = block->FindVar(var_name);
-            // Can't get feed op's TensorDesc
-            if (op_type != "feed" && var_desc && !var_desc->Persistable()) {
-              const auto shape = var_desc->GetShape();
-              if (shape.size() == 1) return false;
-            }
-          }
-        }
-      }
-    }
 
     if (op_type == "dropout") {
       /*
@@ -409,6 +404,7 @@ struct SimpleOpTypeSetTeller : public Teller {
       return false;
 #endif
     }
+
     if (op_type == "softmax") {
       auto* block = desc.Block();
       if (block == nullptr) {
@@ -420,7 +416,17 @@ struct SimpleOpTypeSetTeller : public Teller {
       auto x_var_name = desc.Input("X")[0];
       auto* x_var_desc = block->FindVar(x_var_name);
       const auto x_shape = x_var_desc->GetShape();
+
+      if (with_dynamic_shape && (x_shape.size() == 1 || x_shape.size() == 0)) {
+        int axis = desc.HasAttr("axis")
+                       ? PADDLE_GET_CONST(int, desc.GetAttr("axis"))
+                       : -1;
+        if (axis > 0) {
+          return false;
+        }
+      }
     }
+
     if (op_type == "group_norm") {
       if (!desc.HasAttr("epsilon") || !desc.HasAttr("groups") ||
           !desc.HasAttr("data_layout"))
@@ -977,20 +983,6 @@ struct SimpleOpTypeSetTeller : public Teller {
       }
     }
 
-    if (op_type == "hard_swish") {
-      if (desc.Input("X").size() != 1) {
-        VLOG(3) << "HardSwish op has only 1 input, but got "
-                << desc.Input("X").size();
-        return false;
-      }
-
-      if (desc.Output("Out").size() != 1) {
-        VLOG(3) << "HardSwish op has only 1 output, but got "
-                << desc.Output("Out").size();
-        return false;
-      }
-    }
-
     if (op_type == "squeeze2") {
       // If Attribute is Variable(s), HasAttr() will return False
       if (!desc.HasAttr("axes", /*with_attr_var=*/false)) {
@@ -1204,11 +1196,6 @@ struct SimpleOpTypeSetTeller : public Teller {
               dtype == framework::proto::VarType::FP16)) {
           return false;
         }
-        if (x_shape.size() == 1) {
-          VLOG(3)
-              << "Scale op does not support 1-dimensional input in tensorrt";
-          return false;
-        }
       } else {
         // At present, only support float32 or float16 or int32 into trt.
         if (!(dtype == framework::proto::VarType::FP32 ||
@@ -1513,6 +1500,7 @@ struct SimpleOpTypeSetTeller : public Teller {
                    "elementwise op.";
         return false;
       }
+
       if (x_var_desc->Persistable() && !with_dynamic_shape) {
         VLOG(3)
             << "Input X is a parameter which is not supported for "
@@ -1530,10 +1518,24 @@ struct SimpleOpTypeSetTeller : public Teller {
                "mode.";
         return false;
       }
-    }
-    // remember that 1D input in static shape mode is filtered at the beginning
-    if (op_type == "sum") {
-      return true;
+
+      auto* block = desc.Block();
+      if (block == nullptr) {
+        VLOG(3) << "The block desc is nullptr, we can't continue to analyze. "
+                   "Developers need to check whether block_desc is passed in "
+                   "the pass.";
+        return false;
+      }
+      auto x_var_name = desc.Input("X")[0];
+      auto* x_var_desc = block->FindVar(x_var_name);
+      const auto x_shape = x_var_desc->GetShape();
+      int rank = x_shape.size();
+      int axis = desc.HasAttr("axis")
+                     ? PADDLE_GET_CONST(int, desc.GetAttr("axis"))
+                     : -1;
+      if (axis > rank || axis < -(rank + 1)) {
+        return false;
+      }
     }
 
     if (op_type == "shape" && !with_dynamic_shape) {
@@ -1601,22 +1603,6 @@ struct SimpleOpTypeSetTeller : public Teller {
         if (PADDLE_GET_CONST(bool, desc.GetAttr("approximate"))) return false;
       }
 #endif
-
-      auto* block = desc.Block();
-      if (block == nullptr) {
-        VLOG(3) << "The block desc is nullptr, we can't continue to analyze. "
-                   "Developers need to check whether block_desc is passed in "
-                   "the pass.";
-        return false;
-      }
-
-      auto x_var_name = desc.Input("X")[0];
-      auto* x_var_desc = block->FindVar(x_var_name);
-      const auto x_shape = x_var_desc->GetShape();
-      if (x_shape.size() == 1) {
-        VLOG(3) << "gelu op does not support input's dim is 1 in tensorrt.";
-        return false;
-      }
     }
 
     if (op_type == "layer_norm") {
@@ -1704,20 +1690,6 @@ struct SimpleOpTypeSetTeller : public Teller {
       }
     }
 
-    if (op_type == "leaky_relu") {
-      if (desc.Input("X").size() != 1) {
-        VLOG(3) << "Invalid number of TRT leaky_relu op converter "
-                   "inputs. Expected 1, but received "
-                << desc.Input("X").size();
-        return false;
-      }
-      if (desc.Output("Out").size() != 1) {
-        VLOG(3) << "output of leaky_relu op converter should be 1, got "
-                << desc.Output("Out").size();
-        return false;
-      }
-    }
-
     if (op_type == "pad") {
       if (!desc.HasAttr("pad_value") || !desc.HasAttr("paddings")) return false;
       const float pad_value =
@@ -1786,22 +1758,7 @@ struct SimpleOpTypeSetTeller : public Teller {
         }
       }
     }
-    if (op_type == "swish") {
-      auto* block = desc.Block();
-      if (block == nullptr) {
-        VLOG(3) << "The block desc is nullptr, we can't continue to analyze. "
-                   "Developers need to check whether block_desc is passed in "
-                   "the pass.";
-        return false;
-      }
-      auto x_var_name = desc.Input("X")[0];
-      auto* x_var_desc = block->FindVar(x_var_name);
-      const auto x_shape = x_var_desc->GetShape();
-      if (x_shape.size() == 1) {
-        VLOG(3) << "swish op does not support input's dim is 1 in tensorrt.";
-        return false;
-      }
-    }
+
     if (op_type == "prelu") {
       if (desc.Input("X").size() != 1) {
         VLOG(3) << "Invalid input X's size of prelu TRT converter. "
@@ -1823,28 +1780,18 @@ struct SimpleOpTypeSetTeller : public Teller {
                    "the pass.";
         return false;
       }
-      auto* var_desc = block->FindVar(desc.Input("Alpha")[0]);
-      if (!var_desc) {
+      auto* alpha_var = block->FindVar(desc.Input("Alpha")[0]);
+      if (!alpha_var) {
         VLOG(3) << "Variable Alpha of prelu TRT converter not found.";
         return false;
       }
-
-      auto x_var_name = desc.Input("X")[0];
-      auto* x_var_desc = block->FindVar(x_var_name);
-      const auto x_shape = x_var_desc->GetShape();
-      if (!with_dynamic_shape && x_shape.size() == 1) {
-        VLOG(3) << "prelu op does not support input's dim is 1 in tensorrt "
-                   "with static shape.";
-        return false;
-      }
-
-#if IS_TRT_VERSION_LT(7000)
-      if (!with_dynamic_shape) {
-        // TODO(inference): fix trt6 static plugin error.
-        VLOG(3) << "prelu static plugin in trt6 has bug.";
+      auto alpha_shape = alpha_var->GetShape();
+      if (!with_dynamic_shape && alpha_shape.size() == 0) {
+        VLOG(3) << op_type
+                << " op does not support alpha's dim is 0 in tensorrt "
+                   "static shape mode.";
         return false;
       }
-#endif
     }
 
     if (op_type == "mish") {
@@ -1860,22 +1807,6 @@ struct SimpleOpTypeSetTeller : public Teller {
                 << desc.Output("Out").size() << ".";
         return false;
       }
-
-      auto* block = desc.Block();
-      if (block == nullptr) {
-        VLOG(3) << "The block desc is nullptr, we can't continue to analyze. "
-                   "Developers need to check whether block_desc is passed in "
-                   "the pass.";
-        return false;
-      }
-
-      auto x_var_name = desc.Input("X")[0];
-      auto* x_var_desc = block->FindVar(x_var_name);
-      const auto x_shape = x_var_desc->GetShape();
-      if (x_shape.size() == 1) {
-        VLOG(3) << "mish op does not support input's dim is 1 in tensorrt.";
-        return false;
-      }
     }
 
     if (op_type == "roi_align") {
@@ -2338,26 +2269,6 @@ struct SimpleOpTypeSetTeller : public Teller {
       }
     }
 
-    if (op_type == "hard_sigmoid") {
-      if (!with_dynamic_shape) {
-        auto* block = desc.Block();
-        if (block == nullptr) {
-          VLOG(3) << "The block desc is nullptr, we can't continue to analyze. "
-                     "Developers need to check whether block_desc is passed in "
-                     "the pass.";
-          return false;
-        }
-        auto x_var_name = desc.Input("X")[0];
-        auto* x_var_desc = block->FindVar(x_var_name);
-        const auto x_shape = x_var_desc->GetShape();
-        if (x_shape.size() == 1) {
-          VLOG(3) << "Hard sigmoid does not support 1-dimensional input in "
-                     "tensorrt";
-          return false;
-        }
-      }
-    }
-
     if (op_type == "cast") {
 // trt 6015 result in Windows ppyolo_mbv3 TRT fp32 diff
 #if !IS_TRT_VERSION_GE(7000)
@@ -2592,6 +2503,15 @@ struct SimpleOpTypeSetTeller : public Teller {
                    "the pass.";
         return false;
       }
+
+#if IS_TRT_VERSION_LT(8000)
+      auto x_var_name = desc.Input("X")[0];
+      auto* x_var_desc = block->FindVar(x_var_name);
+      const auto x_shape = x_var_desc->GetShape();
+      if (x_shape.size() == 0) {
+        return false;  // not supported 0 dim.
+      }
+#endif
     }
 
     if (op_type == "grid_sampler") {
@@ -2908,6 +2828,7 @@ struct SimpleOpTypeSetTeller : public Teller {
       "elementwise_mul",
       "elementwise_div",
       "elementwise_pow",
+      "pow",
       "elementwise_min",
       "elementwise_max",
       "elementwise_floordiv",

paddle/phi/infermeta/binary.cc

@@ -2328,7 +2328,7 @@ void PReluInferMeta(const MetaTensor& x,
                       1,
                       phi::errors::InvalidArgument(
                           "For mode 'element', rank of input X must be "
-                          "equal or larger than 2. But recevied X's "
+                          "equal or larger than 1. But recevied X's "
                           "rank: %d",
                           x_rank));
     PADDLE_ENFORCE_EQ(

test/ir/inference/test_trt_convert_activation.py

@@ -60,6 +60,9 @@ class TrtConvertActivationTest(TrtLayerAutoScanTest):
                     "logsigmoid",
                     "tanh_shrink",
                     "softplus",
+                    "hard_swish",
+                    "hard_sigmoid",
+                    "leaky_relu",
                 ]:
                     # few samples to reduce time
                     # for beta in [-0.2, 0.5, 0.67, 3]:
@@ -80,6 +83,18 @@ class TrtConvertActivationTest(TrtLayerAutoScanTest):
                                 dics = [{"threshold": alpha}]
                             if op_type == "softplus":
                                 dics = [{"beta": beta}]
+                            if op_type == "hard_swish":
+                                dics = [
+                                    {
+                                        "threshold": 6.0,
+                                        "scale": 6.0,
+                                        "offset": 3.0,
+                                    }
+                                ]
+                            if op_type == "hard_sigmoid":
+                                dics = [{"slope": beta, "offset": alpha}]
+                            if op_type == "leaky_relu":
+                                dics = [{"alpha": alpha}]
 
                             ops_config = [
                                 {

test/ir/inference/test_trt_convert_elementwise.py

@@ -1092,5 +1092,283 @@ class TrtConvertElementwiseTestTwoInputSkipCase(TrtLayerAutoScanTest):
         self.run_test()
 
 
+class TrtConvertPowOp(TrtLayerAutoScanTest):
+    def is_program_valid(self, program_config: ProgramConfig) -> bool:
+        return True
+
+    def sample_program_configs(self):
+        def generate_input(shape):
+            if len(shape) == 0:
+                return np.random.random([]).astype(np.float32)
+            return np.random.random(shape).astype(np.float32)
+
+        for batch in [1, 4]:
+            for shape in [
+                [],
+                [32],
+                [batch, 32],
+                [batch, 32, 32],
+                [batch, 32, 16, 32],
+            ]:
+                for factor in [1.0, 2.0, -1.0, 0.5, -2]:
+                    self.dims = len(shape)
+                    dics = [{"factor": factor}]
+                    ops_config = [
+                        {
+                            "op_type": "pow",
+                            "op_inputs": {
+                                "X": ["input_data"],
+                            },
+                            "op_outputs": {"Out": ["output_data"]},
+                            "op_attrs": dics[0],
+                            "outputs_dtype": {"output_data": np.float32},
+                        }
+                    ]
+                    ops = self.generate_op_config(ops_config)
+
+                    program_config = ProgramConfig(
+                        ops=ops,
+                        weights={},
+                        inputs={
+                            "input_data": TensorConfig(
+                                data_gen=partial(generate_input, shape)
+                            ),
+                        },
+                        outputs=["output_data"],
+                    )
+
+                    yield program_config
+
+    def sample_predictor_configs(
+        self, program_config
+    ) -> (paddle_infer.Config, List[int], float):
+        def generate_dynamic_shape(attrs):
+            if self.dims == 0:
+                self.dynamic_shape.min_input_shape = {"input_data": []}
+                self.dynamic_shape.max_input_shape = {"input_data": []}
+                self.dynamic_shape.opt_input_shape = {"input_data": []}
+            elif self.dims == 1:
+                self.dynamic_shape.min_input_shape = {"input_data": [4]}
+                self.dynamic_shape.max_input_shape = {"input_data": [32]}
+                self.dynamic_shape.opt_input_shape = {"input_data": [16]}
+            elif self.dims == 2:
+                self.dynamic_shape.min_input_shape = {"input_data": [1, 32]}
+                self.dynamic_shape.max_input_shape = {"input_data": [4, 32]}
+                self.dynamic_shape.opt_input_shape = {"input_data": [2, 32]}
+            elif self.dims == 3:
+                self.dynamic_shape.min_input_shape = {"input_data": [1, 32, 4]}
+                self.dynamic_shape.max_input_shape = {"input_data": [4, 32, 32]}
+                self.dynamic_shape.opt_input_shape = {"input_data": [2, 32, 32]}
+            elif self.dims == 4:
+                self.dynamic_shape.min_input_shape = {
+                    "input_data": [1, 32, 4, 4]
+                }
+                self.dynamic_shape.max_input_shape = {
+                    "input_data": [4, 32, 32, 32]
+                }
+                self.dynamic_shape.opt_input_shape = {
+                    "input_data": [4, 32, 16, 32]
+                }
+
+        def clear_dynamic_shape():
+            self.dynamic_shape.max_input_shape = {}
+            self.dynamic_shape.min_input_shape = {}
+            self.dynamic_shape.opt_input_shape = {}
+
+        def generate_trt_nodes_num(attrs, dynamic_shape):
+            if (self.dims == 1 or self.dims == 0) and not dynamic_shape:
+                return 0, 3
+            return 1, 2
+
+        attrs = [
+            program_config.ops[i].attrs for i in range(len(program_config.ops))
+        ]
+
+        # for static_shape
+        clear_dynamic_shape()
+        self.trt_param.precision = paddle_infer.PrecisionType.Float32
+        yield self.create_inference_config(), generate_trt_nodes_num(
+            attrs, False
+        ), (1e-5, 1e-5)
+        self.trt_param.precision = paddle_infer.PrecisionType.Half
+        yield self.create_inference_config(), generate_trt_nodes_num(
+            attrs, False
+        ), (1e-3, 1e-3)
+
+        # for dynamic_shape
+        generate_dynamic_shape(attrs)
+        self.trt_param.precision = paddle_infer.PrecisionType.Float32
+        yield self.create_inference_config(), generate_trt_nodes_num(
+            attrs, True
+        ), (1e-5, 1e-5)
+        self.trt_param.precision = paddle_infer.PrecisionType.Half
+        yield self.create_inference_config(), generate_trt_nodes_num(
+            attrs, True
+        ), (1e-3, 1e-3)
+
+    def add_skip_trt_case(self):
+        pass
+
+    def test(self):
+        self.add_skip_trt_case()
+        self.run_test()
+
+
+class TrtConvertElementwise0D(TrtLayerAutoScanTest):
+    def is_program_valid(self, program_config: ProgramConfig) -> bool:
+        return True
+
+    def sample_program_configs(self):
+        def generate_input(dims, op_type):
+            shape = []
+            if dims == 0:
+                shape = []
+            elif dims == 1:
+                shape = [8]
+            elif dims == 2:
+                shape = [1, 8]
+            elif dims == 3:
+                shape = [1, 8, 8]
+            else:
+                shape = [1, 8, 8, 8]
+
+            # elementwise_floordiv is integer only
+            if op_type == "elementwise_floordiv":
+                return np.random.randint(
+                    low=1, high=10000, size=shape, dtype=np.int32
+                )
+            elif op_type == "elementwise_mod":
+                return np.random.uniform(low=0.1, high=1.0, size=shape).astype(
+                    np.float32
+                )
+            else:
+                return np.random.random(shape).astype(np.float32)
+
+        for dims in [[0, 0], [0, 1], [0, 2], [1, 0], [2, 0]]:
+            for op_type in [
+                "elementwise_add",
+                "elementwise_mul",
+                "elementwise_sub",
+                "elementwise_div",
+                "elementwise_pow",
+                "elementwise_min",
+                "elementwise_max",
+                "elementwise_floordiv",
+                "elementwise_mod",
+            ]:
+                for axis in [-1 if dims[0] == 1 or dims[0] == 0 else 1]:
+                    self.dims = dims[0]
+                    dics = [{"axis": axis}]
+                    ops_config = [
+                        {
+                            "op_type": op_type,
+                            "op_inputs": {
+                                "X": ["input_data"],
+                                "Y": ["weight"],
+                            },
+                            "op_outputs": {"Out": ["output_data"]},
+                            "op_attrs": dics[0],
+                            "outputs_dtype": {
+                                "output_data": np.float32
+                                if op_type != "elementwise_floordiv"
+                                else np.int32
+                            },
+                        }
+                    ]
+                    ops = self.generate_op_config(ops_config)
+
+                    program_config = ProgramConfig(
+                        ops=ops,
+                        weights={
+                            "weight": TensorConfig(
+                                data_gen=partial(
+                                    generate_input, dims[1], op_type
+                                )
+                            )
+                        },
+                        inputs={
+                            "input_data": TensorConfig(
+                                data_gen=partial(
+                                    generate_input, dims[0], op_type
+                                )
+                            ),
+                        },
+                        outputs=["output_data"],
+                    )
+
+                    yield program_config
+
+    def sample_predictor_configs(
+        self, program_config
+    ) -> (paddle_infer.Config, List[int], float):
+        def generate_dynamic_shape(attrs):
+            # The input.dims[1] must be equal to the weight's length.
+            if self.dims == 0:
+                self.dynamic_shape.min_input_shape = {"input_data": []}
+                self.dynamic_shape.max_input_shape = {"input_data": []}
+                self.dynamic_shape.opt_input_shape = {"input_data": []}
+            if self.dims == 1:
+                self.dynamic_shape.min_input_shape = {"input_data": [1]}
+                self.dynamic_shape.max_input_shape = {"input_data": [16]}
+                self.dynamic_shape.opt_input_shape = {"input_data": [8]}
+            elif self.dims == 2:
+                self.dynamic_shape.min_input_shape = {"input_data": [1, 8]}
+                self.dynamic_shape.max_input_shape = {"input_data": [4, 8]}
+                self.dynamic_shape.opt_input_shape = {"input_data": [2, 8]}
+            elif self.dims == 3:
+                self.dynamic_shape.min_input_shape = {"input_data": [1, 1, 4]}
+                self.dynamic_shape.max_input_shape = {"input_data": [4, 16, 16]}
+                self.dynamic_shape.opt_input_shape = {"input_data": [2, 8, 8]}
+            elif self.dims == 4:
+                self.dynamic_shape.min_input_shape = {
+                    "input_data": [1, 8, 8, 8]
+                }
+                self.dynamic_shape.max_input_shape = {
+                    "input_data": [4, 8, 8, 8]
+                }
+                self.dynamic_shape.opt_input_shape = {
+                    "input_data": [4, 8, 8, 8]
+                }
+
+        def clear_dynamic_shape():
+            self.dynamic_shape.max_input_shape = {}
+            self.dynamic_shape.min_input_shape = {}
+            self.dynamic_shape.opt_input_shape = {}
+
+        def generate_trt_nodes_num(attrs, dynamic_shape):
+            if not dynamic_shape and (self.dims == 1 or self.dims == 0):
+                return 0, 3
+            return 1, 2
+
+        attrs = [
+            program_config.ops[i].attrs for i in range(len(program_config.ops))
+        ]
+
+        # for static_shape
+        clear_dynamic_shape()
+        self.trt_param.precision = paddle_infer.PrecisionType.Float32
+        yield self.create_inference_config(), generate_trt_nodes_num(
+            attrs, False
+        ), (1e-5, 1e-5)
+        self.trt_param.precision = paddle_infer.PrecisionType.Half
+        yield self.create_inference_config(), generate_trt_nodes_num(
+            attrs, False
+        ), (1e-3, 1e-3)
+
+        # # for dynamic_shape
+        generate_dynamic_shape(attrs)
+        self.trt_param.precision = paddle_infer.PrecisionType.Float32
+        yield self.create_inference_config(), generate_trt_nodes_num(
+            attrs, True
+        ), (1e-5, 1e-5)
+        self.trt_param.precision = paddle_infer.PrecisionType.Half
+        yield self.create_inference_config(), generate_trt_nodes_num(
+            attrs, True
+        ), (1e-3, 1e-3)
+
+    def test(self):
+        self.run_test()
+
+
 if __name__ == "__main__":
     unittest.main()

test/ir/inference/test_trt_convert_equal.py

@@ -40,54 +40,64 @@ class TrtConvertEqualOneInputCornerCase(TrtLayerAutoScanTest):
             return np.random.random(shape).astype(np.float32)
 
         for op_type in ["equal", "not_equal"]:
-            for batch in [1, 2, 4]:
-                for shape in [[batch, 1], [batch, 1, 32], [batch, 1, 16, 32]]:
-                    for axis in [-1 if len(shape) == 1 else 1]:
-                        self.dims = len(shape)
-                        dics = [{"axis": axis}, {"in_dtype": 0, "out_dtype": 5}]
-                        ops_config = [
-                            {
-                                "op_type": op_type,
-                                "op_inputs": {
-                                    "X": ["input_data1"],
-                                    "Y": ["input_data2"],
-                                },
-                                "op_outputs": {"Out": ["compare_output_data"]},
-                                "op_attrs": dics[0],
-                                "outputs_dtype": {
-                                    "compare_output_data": np.bool_
-                                },
+            for shape in [[], [1, 1], [1, 1, 32], [1, 1, 16, 32]]:
+                for axis in [-1 if len(shape) == 1 or len(shape) == 0 else 1]:
+                    self.dims = len(shape)
+                    dics = [{"axis": axis}, {"in_dtype": 0, "out_dtype": 5}]
+                    ops_config = [
+                        {
+                            "op_type": op_type,
+                            "op_inputs": {
+                                "X": ["input_data1"],
+                                "Y": ["input_data2"],
                             },
-                            {
-                                "op_type": "cast",
-                                "op_inputs": {"X": ["compare_output_data"]},
-                                "op_outputs": {"Out": ["output_data"]},
-                                "op_attrs": dics[1],
-                                "outputs_dtype": {"output_data": np.float32},
-                            },
-                        ]
-                        ops = self.generate_op_config(ops_config)
-
-                        program_config = ProgramConfig(
-                            ops=ops,
-                            weights={},
-                            inputs={
-                                "input_data1": TensorConfig(
-                                    data_gen=partial(generate_input, shape)
-                                ),
-                                "input_data2": TensorConfig(
-                                    data_gen=partial(generate_input, shape)
-                                ),
-                            },
-                            outputs=["output_data"],
-                        )
-                        yield program_config
+                            "op_outputs": {"Out": ["compare_output_data"]},
+                            "op_attrs": dics[0],
+                            "outputs_dtype": {"compare_output_data": np.bool_},
+                        },
+                        {
+                            "op_type": "cast",
+                            "op_inputs": {"X": ["compare_output_data"]},
+                            "op_outputs": {"Out": ["output_data"]},
+                            "op_attrs": dics[1],
+                            "outputs_dtype": {"output_data": np.float32},
+                        },
+                    ]
+                    ops = self.generate_op_config(ops_config)
+
+                    program_config = ProgramConfig(
+                        ops=ops,
+                        weights={},
+                        inputs={
+                            "input_data1": TensorConfig(
+                                data_gen=partial(generate_input, shape)
+                            ),
+                            "input_data2": TensorConfig(
+                                data_gen=partial(generate_input, shape)
+                            ),
+                        },
+                        outputs=["output_data"],
+                    )
+                    yield program_config
 
     def sample_predictor_configs(
         self, program_config
     ) -> (paddle_infer.Config, List[int], float):
         def generate_dynamic_shape(attrs):
             # The input.dims[1] must be equal to the weight's length.
+            if self.dims == 0:
+                self.dynamic_shape.min_input_shape = {
+                    "input_data1": [],
+                    "input_data2": [],
+                }
+                self.dynamic_shape.max_input_shape = {
+                    "input_data1": [],
+                    "input_data2": [],
+                }
+                self.dynamic_shape.opt_input_shape = {
+                    "input_data1": [],
+                    "input_data2": [],
+                }
             if self.dims == 2:
                 self.dynamic_shape.min_input_shape = {
                     "input_data1": [1, 1],

test/ir/inference/test_trt_convert_expand_as_v2.py

@@ -49,8 +49,11 @@ class TrtConvertExpandASV2Test(TrtLayerAutoScanTest):
             elif self.dims == 1:
                 self.input_shape = [32]
                 return np.random.random([32]).astype(np.float32)
+            elif self.dims == 0:
+                self.input_shape = []
+                return np.random.random([]).astype(np.float32)
 
-        for dims in [1, 2, 3, 4]:
+        for dims in [0, 1, 2, 3, 4]:
             for shape in [
                 [10, 8, 32, 32],
                 [2, 8, 32, 32],
@@ -125,6 +128,10 @@ class TrtConvertExpandASV2Test(TrtLayerAutoScanTest):
                 self.dynamic_shape.min_input_shape = {"expand_v2_input": [32]}
                 self.dynamic_shape.max_input_shape = {"expand_v2_input": [64]}
                 self.dynamic_shape.opt_input_shape = {"expand_v2_input": [32]}
+            elif self.dims == 0:
+                self.dynamic_shape.min_input_shape = {"expand_v2_input": []}
+                self.dynamic_shape.max_input_shape = {"expand_v2_input": []}
+                self.dynamic_shape.opt_input_shape = {"expand_v2_input": []}
 
         def clear_dynamic_shape():
             self.dynamic_shape.min_input_shape = {}
@@ -132,7 +139,9 @@ class TrtConvertExpandASV2Test(TrtLayerAutoScanTest):
             self.dynamic_shape.opt_input_shape = {}
 
         def generate_trt_nodes_num(attrs, dynamic_shape):
-            if dynamic_shape:
+            ver = paddle_infer.get_trt_compile_version()
+            ver_num = ver[0] * 1000 + ver[1] * 100 + ver[2] * 10
+            if dynamic_shape and (ver_num > 8000 or self.dims > 0):
                 return 1, 2
             else:
                 return 0, 3

test/ir/inference/test_trt_convert_gelu.py

@@ -29,7 +29,9 @@ class TrtConvertGeluTest(TrtLayerAutoScanTest):
 
     def sample_program_configs(self):
         def generate_input1(dims, attrs: List[Dict[str, Any]]):
-            if dims == 1:
+            if dims == 0:
+                return np.ones([]).astype(np.float32)
+            elif dims == 1:
                 return np.ones([32]).astype(np.float32)
             elif dims == 2:
                 return np.ones([3, 32]).astype(np.float32)
@@ -38,7 +40,7 @@ class TrtConvertGeluTest(TrtLayerAutoScanTest):
             else:
                 return np.ones([1, 3, 32, 32]).astype(np.float32)
 
-        for dims in [1, 2, 3, 4]:
+        for dims in [0, 1, 2, 3, 4]:
             for approximate in [True, False]:
                 self.dims = dims
                 dics = [{"approximate": approximate}]
@@ -70,7 +72,11 @@ class TrtConvertGeluTest(TrtLayerAutoScanTest):
         self, program_config
     ) -> (paddle_infer.Config, List[int], float):
         def generate_dynamic_shape(attrs):
-            if self.dims == 1:
+            if self.dims == 0:
+                self.dynamic_shape.min_input_shape = {"input_data": []}
+                self.dynamic_shape.max_input_shape = {"input_data": []}
+                self.dynamic_shape.opt_input_shape = {"input_data": []}
+            elif self.dims == 1:
                 self.dynamic_shape.min_input_shape = {"input_data": [1]}
                 self.dynamic_shape.max_input_shape = {"input_data": [64]}
                 self.dynamic_shape.opt_input_shape = {"input_data": [32]}
@@ -104,7 +110,7 @@ class TrtConvertGeluTest(TrtLayerAutoScanTest):
             runtime_version = paddle_infer.get_trt_runtime_version()
             self.assertTrue(compile_version == runtime_version)
             # Dimension one only runs on Paddle OP
-            if self.dims == 1:
+            if not dynamic_shape and (self.dims == 1 or self.dims == 0):
                 return 0, 3
             if compile_version >= valid_version:
                 return 1, 2

test/ir/inference/test_trt_convert_prelu.py

@@ -18,7 +18,7 @@ from typing import Any, Dict, List
 
 import numpy as np
 from program_config import ProgramConfig, TensorConfig
-from trt_layer_auto_scan_test import SkipReasons, TrtLayerAutoScanTest
+from trt_layer_auto_scan_test import TrtLayerAutoScanTest
 
 import paddle.inference as paddle_infer
 
@@ -28,170 +28,165 @@ class TrtConvertPreluTest(TrtLayerAutoScanTest):
         return True
 
     def sample_program_configs(self):
-        def generate_input(batch, dim1, dim2, dim3):
-            shape = [batch]
-            if dim1 != 0:
-                shape.append(dim1)
-            if dim2 != 0:
-                shape.append(dim2)
-            if dim3 != 0:
-                shape.append(dim3)
-            return np.random.random(shape).astype(np.float32)
-
-        def generate_alpha(attrs: List[Dict[str, Any]], dim1, dim2, dim3):
+        def generate_input(attrs: List[Dict[str, Any]], batch):
+            if self.dims == 0:
+                return np.random.random([]).astype(np.float32)
+            elif self.dims == 1:
+                return np.random.random([16]).astype(np.float32)
+            elif self.dims == 2:
+                return np.random.random([1, 3]).astype(np.float32)
+            elif self.dims == 3:
+                if attrs[0]["data_format"] == "NCHW":
+                    return np.random.random([batch, 3, 16]).astype(np.float32)
+                elif attrs[0]["data_format"] == "NHWC":
+                    return np.random.random([batch, 16, 3]).astype(np.float32)
+                else:
+                    raise AssertionError()
+            else:
+                if attrs[0]["data_format"] == "NCHW":
+                    return np.random.random([batch, 3, 16, 32]).astype(
+                        np.float32
+                    )
+                else:
+                    return np.random.random([batch, 16, 32, 3]).astype(
+                        np.float32
+                    )
+
+        def generate_alpha(attrs: List[Dict[str, Any]]):
+            if self.dims == 0:
+                return np.random.random([]).astype(np.float32)
             if attrs[0]["mode"] == "all":
-                return np.random.random(size=(1)).astype(np.float32)
-            elif (
-                attrs[0]["mode"] == "channel"
-                and attrs[0]["data_format"] == "NCHW"
-            ):
-                shape = [1]
-                if dim1 != 0:
-                    shape.append(dim1)
-                if dim2 != 0:
-                    shape.append(dim2)
-                if dim3 != 0:
-                    shape.append(dim3)
-                return np.random.random(size=shape[1]).astype(np.float32)
-            elif (
-                attrs[0]["mode"] == "channel"
-                and attrs[0]["data_format"] == "NHWC"
-            ):
-                shape = [1]
-                if dim1 != 0:
-                    shape.append(dim1)
-                if dim2 != 0:
-                    shape.append(dim2)
-                if dim3 != 0:
-                    shape.append(dim3)
-                return np.random.random(size=shape[-1]).astype(np.float32)
+                return np.random.random([1]).astype(np.float32)
+            elif attrs[0]["mode"] == "channel":
+                return np.random.random([3]).astype(np.float32)
             elif attrs[0]["mode"] == "element":
-                shape = [1]
-                if dim1 != 0:
-                    shape.append(dim1)
-                if dim2 != 0:
-                    shape.append(dim2)
-                if dim3 != 0:
-                    shape.append(dim3)
-                return np.random.random(size=shape).astype(np.float32)
+                if self.dims == 1:
+                    return np.random.random([16]).astype(np.float32)
+                elif self.dims == 2:
+                    return np.random.random([1, 3]).astype(np.float32)
+                elif self.dims == 3:
+                    if attrs[0]["data_format"] == "NCHW":
+                        return np.random.random([1, 3, 16]).astype(np.float32)
+                    elif attrs[0]["data_format"] == "NHWC":
+                        return np.random.random([1, 16, 3]).astype(np.float32)
+                    else:
+                        raise AssertionError()
+                else:
+                    if attrs[0]["data_format"] == "NCHW":
+                        return np.random.random([1, 3, 16, 32]).astype(
+                            np.float32
+                        )
+                    elif attrs[0]["data_format"] == "NHWC":
+                        return np.random.random([1, 16, 32, 3]).astype(
+                            np.float32
+                        )
+                    else:
+                        raise AssertionError()
 
         for batch in [1, 4]:
-            for dim1 in [0, 3]:
-                for dim2 in [0, 16]:
-                    for dim3 in [0, 32]:
-                        self.dim1 = dim1
-                        self.dim2 = dim2
-                        self.dim3 = dim3
-
-                        if dim1 == 0 and dim2 != 0:
+            for dims in [0, 1, 2, 3, 4]:
+                for mode in ["all", "element", "channel"]:
+                    for data_format in ["NCHW", "NHWC"]:
+                        if (mode == "element" or mode == "all") and dims == 0:
                             continue
-                        if dim1 == 0 and dim2 == 0 and dim3 != 0:
+                        if mode == "channel" and dims != 4:
                             continue
-
-                        for mode in ["all", "channel", "element"]:
-                            for data_format in ['NCHW', 'NHWC']:
-                                if (
-                                    mode == "channel"
-                                    and dim1 == 0
-                                    and data_format == "NCHW"
-                                ):
-                                    continue
-                                if (
-                                    mode == "channel"
-                                    and dim3 == 0
-                                    and data_format == "NHWC"
-                                ):
-                                    continue
-                                dics = [
-                                    {"mode": mode, "data_format": data_format}
-                                ]
-                                ops_config = [
-                                    {
-                                        "op_type": "prelu",
-                                        "op_inputs": {
-                                            "X": ["input_data"],
-                                            "Alpha": ["alpha_weight"],
-                                        },
-                                        "op_outputs": {"Out": ["output_data"]},
-                                        "op_attrs": dics[0],
-                                    }
-                                ]
-                                ops = self.generate_op_config(ops_config)
-
-                                program_config = ProgramConfig(
-                                    ops=ops,
-                                    weights={
-                                        "alpha_weight": TensorConfig(
-                                            data_gen=partial(
-                                                generate_alpha,
-                                                dics,
-                                                dim1,
-                                                dim2,
-                                                dim3,
-                                            )
-                                        )
-                                    },
-                                    inputs={
-                                        "input_data": TensorConfig(
-                                            data_gen=partial(
-                                                generate_input,
-                                                batch,
-                                                dim1,
-                                                dim2,
-                                                dim3,
-                                            )
-                                        ),
-                                    },
-                                    outputs=["output_data"],
+                        self.dims = dims
+                        dics = [{"mode": mode, "data_format": data_format}]
+                        ops_config = [
+                            {
+                                "op_type": "prelu",
+                                "op_inputs": {
+                                    "X": ["input_data"],
+                                    "Alpha": ["alpha_weight"],
+                                },
+                                "op_outputs": {"Out": ["output_data"]},
+                                "op_attrs": dics[0],
+                            }
+                        ]
+                        ops = self.generate_op_config(ops_config)
+
+                        program_config = ProgramConfig(
+                            ops=ops,
+                            weights={
+                                "alpha_weight": TensorConfig(
+                                    data_gen=partial(generate_alpha, dics)
                                 )
-
-                                yield program_config
+                            },
+                            inputs={
+                                "input_data": TensorConfig(
+                                    data_gen=partial(
+                                        generate_input, dics, batch
+                                    )
+                                ),
+                            },
+                            outputs=["output_data"],
+                        )
+
+                        yield program_config
 
     def sample_predictor_configs(
         self, program_config
     ) -> (paddle_infer.Config, List[int], float):
         def generate_dynamic_shape(attrs):
-            if self.dim1 == 0:
-                self.dynamic_shape.min_input_shape = {
-                    "input_data": [1],
-                }
-                self.dynamic_shape.max_input_shape = {
-                    "input_data": [4],
-                }
-                self.dynamic_shape.opt_input_shape = {
-                    "input_data": [2],
-                }
-            else:
-                if self.dim2 == 0 and self.dim3 == 0:
+            if self.dims == 0:
+                self.dynamic_shape.min_input_shape = {"input_data": []}
+                self.dynamic_shape.max_input_shape = {"input_data": []}
+                self.dynamic_shape.opt_input_shape = {"input_data": []}
+            elif self.dims == 1:
+                self.dynamic_shape.min_input_shape = {"input_data": [16]}
+                self.dynamic_shape.max_input_shape = {"input_data": [16]}
+                self.dynamic_shape.opt_input_shape = {"input_data": [16]}
+            elif self.dims == 2:
+                self.dynamic_shape.min_input_shape = {"input_data": [1, 3]}
+                self.dynamic_shape.max_input_shape = {"input_data": [1, 3]}
+                self.dynamic_shape.opt_input_shape = {"input_data": [1, 3]}
+            elif self.dims == 3:
+                if attrs[0]["data_format"] == "NCHW":
                     self.dynamic_shape.min_input_shape = {
-                        "input_data": [1, 1],
+                        "input_data": [1, 3, 16]
                     }
                     self.dynamic_shape.max_input_shape = {
-                        "input_data": [4, 32],
+                        "input_data": [4, 3, 16]
                     }
                     self.dynamic_shape.opt_input_shape = {
-                        "input_data": [2, 3],
+                        "input_data": [1, 3, 16]
                     }
-                elif self.dim2 != 0 and self.dim3 != 0:
+                elif attrs[0]["data_format"] == "NHWC":
                     self.dynamic_shape.min_input_shape = {
-                        "input_data": [1, 1, 1, 1],
+                        "input_data": [1, 16, 3]
                     }
                     self.dynamic_shape.max_input_shape = {
-                        "input_data": [4, 3, 16, 32],
+                        "input_data": [4, 16, 3]
                     }
                     self.dynamic_shape.opt_input_shape = {
-                        "input_data": [2, 3, 16, 32],
+                        "input_data": [1, 16, 3]
                     }
-                elif self.dim3 == 0:
+                else:
+                    raise AssertionError()
+            else:
+                if attrs[0]["data_format"] == "NCHW":
                     self.dynamic_shape.min_input_shape = {
-                        "input_data": [1, 1, 1],
+                        "input_data": [1, 3, 16, 32]
                     }
                     self.dynamic_shape.max_input_shape = {
-                        "input_data": [4, 3, 32],
+                        "input_data": [4, 3, 16, 32]
                     }
                     self.dynamic_shape.opt_input_shape = {
-                        "input_data": [2, 3, 16],
+                        "input_data": [1, 3, 16, 32]
                     }
+                elif attrs[0]["data_format"] == "NHWC":
+                    self.dynamic_shape.min_input_shape = {
+                        "input_data": [1, 16, 32, 3]
+                    }
+                    self.dynamic_shape.max_input_shape = {
+                        "input_data": [4, 16, 32, 3]
+                    }
+                    self.dynamic_shape.opt_input_shape = {
+                        "input_data": [1, 16, 32, 3]
+                    }
+                else:
+                    raise AssertionError()
 
         def clear_dynamic_shape():
             self.dynamic_shape.max_input_shape = {}
@@ -203,12 +198,7 @@ class TrtConvertPreluTest(TrtLayerAutoScanTest):
         ]
 
         def generate_trt_nodes_num(attrs, dynamic_shape):
-            if (
-                not dynamic_shape
-                and self.dim1 == 0
-                and self.dim2 == 0
-                and self.dim3 == 0
-            ):
+            if not dynamic_shape and (self.dims == 1 or self.dims == 0):
                 return 0, 3
             return 1, 2
 
@@ -234,23 +224,7 @@ class TrtConvertPreluTest(TrtLayerAutoScanTest):
             attrs, True
         ), (1e-3, 1e-3)
 
-    def add_skip_trt_case(self):
-        ver = paddle_infer.get_trt_compile_version()
-        if ver[0] * 1000 + ver[1] * 100 + ver[0] * 10 < 7000:
-
-            def teller(program_config, predictor_config):
-                if not predictor_config.tensorrt_dynamic_shape_enabled():
-                    return True
-                return False
-
-            self.add_skip_case(
-                teller,
-                SkipReasons.TRT_NOT_IMPLEMENTED,
-                "Need to repair the case: the output of GPU and tensorrt has diff in trt6, the prelu static plugin has bug.",
-            )
-
     def test(self):
-        self.add_skip_trt_case()
         self.run_test()
 
 

test/ir/inference/test_trt_convert_reshape.py

@@ -431,5 +431,99 @@ class TrtConvertReshapeTest3(TrtLayerAutoScanTest):
         self.run_test()
 
 
+class TrtConvertReshapeZeroDimsTest(TrtLayerAutoScanTest):
+    def is_program_valid(self, program_config: ProgramConfig) -> bool:
+        return True
+
+    def sample_program_configs(self):
+        def generate_input1(attrs: List[Dict[str, Any]]):
+            if self.dims > 0:
+                self.input_shape = [1] * self.dims
+                return np.random.random(self.input_shape).astype(np.float32)
+            elif self.dims == 0:
+                self.input_shape = []
+                return np.random.random([]).astype(np.float32)
+
+        for dims in [0, 1, 2, 3]:
+            for shape in [
+                [],
+                [1, 1],
+            ]:
+                dics = [
+                    {
+                        "shape": shape,
+                    },
+                ]
+                self.dims = dims
+                dics_intput = [{"X": ["reshape_input"]}]
+
+                ops_config = [
+                    {
+                        "op_type": "reshape",
+                        "op_inputs": dics_intput[0],
+                        "op_outputs": {"Out": ["reshape_out"]},
+                        "op_attrs": dics[0],
+                    }
+                ]
+                ops = self.generate_op_config(ops_config)
+                program_config = ProgramConfig(
+                    ops=ops,
+                    weights={},
+                    inputs={
+                        "reshape_input": TensorConfig(
+                            data_gen=partial(generate_input1, dics)
+                        )
+                    },
+                    outputs=["reshape_out"],
+                )
+
+                yield program_config
+
+    def sample_predictor_configs(
+        self, program_config
+    ) -> (paddle_infer.Config, List[int], float):
+        def generate_dynamic_shape(attrs):
+            self.dynamic_shape.min_input_shape = {
+                "reshape_input": self.input_shape
+            }
+            self.dynamic_shape.max_input_shape = {
+                "reshape_input": self.input_shape
+            }
+            self.dynamic_shape.opt_input_shape = {
+                "reshape_input": self.input_shape
+            }
+
+        def clear_dynamic_shape():
+            self.dynamic_shape.min_input_shape = {}
+            self.dynamic_shape.max_input_shape = {}
+            self.dynamic_shape.opt_input_shape = {}
+
+        def generate_trt_nodes_num(attrs, dynamic_shape):
+            # only test dynamic shape mode
+            return 1, 2
+
+        attrs = [
+            program_config.ops[i].attrs for i in range(len(program_config.ops))
+        ]
+
+        # for dynamic_shape
+        generate_dynamic_shape(attrs)
+        self.trt_param.precision = paddle_infer.PrecisionType.Float32
+        yield self.create_inference_config(), generate_trt_nodes_num(
+            attrs, True
+        ), 1e-5
+        self.trt_param.precision = paddle_infer.PrecisionType.Half
+        yield self.create_inference_config(), generate_trt_nodes_num(
+            attrs, True
+        ), 1e-3
+
+    def add_skip_trt_case(self):
+        pass
+
+    def test(self):
+        self.add_skip_trt_case()
+        self.run_test()
+
+
 if __name__ == "__main__":
     unittest.main()

test/ir/inference/test_trt_convert_scale.py

@@ -43,12 +43,14 @@ class TrtConvertScaleTest(TrtLayerAutoScanTest):
                 )
             elif self.dims == 1:
                 return np.ones([24]).astype(np.int32 if is_int else np.float32)
+            elif self.dims == 0:
+                return np.ones([]).astype(np.int32 if is_int else np.float32)
 
         def generate_weight1(attrs: List[Dict[str, Any]], is_int):
             return np.ones([1]).astype(np.int32 if is_int else np.float32)
 
         for num_input in [0, 1]:
-            for dims in [1, 2, 3, 4]:
+            for dims in [0, 1, 2, 3, 4]:
                 for batch in [1, 2]:
                     for scale in [0.1, -1.0]:
                         for bias in [0.0, 1.2]:
@@ -141,6 +143,10 @@ class TrtConvertScaleTest(TrtLayerAutoScanTest):
                 self.dynamic_shape.min_input_shape = {"scale_input": [24]}
                 self.dynamic_shape.max_input_shape = {"scale_input": [48]}
                 self.dynamic_shape.opt_input_shape = {"scale_input": [24]}
+            elif self.dims == 0:
+                self.dynamic_shape.min_input_shape = {"scale_input": []}
+                self.dynamic_shape.max_input_shape = {"scale_input": []}
+                self.dynamic_shape.opt_input_shape = {"scale_input": []}
 
         def clear_dynamic_shape():
             self.dynamic_shape.min_input_shape = {}
@@ -148,6 +154,8 @@ class TrtConvertScaleTest(TrtLayerAutoScanTest):
             self.dynamic_shape.opt_input_shape = {}
 
         def generate_trt_nodes_num(attrs, dynamic_shape):
+            if not dynamic_shape and (self.dims == 1 or self.dims == 0):
+                return 0, 3
             return 1, 2
 
         attrs = [
@@ -189,23 +197,12 @@ class TrtConvertScaleTest(TrtLayerAutoScanTest):
         )
 
         def teller2(program_config, predictor_config):
-            if self.dims == 1 and len(self.dynamic_shape.min_input_shape) == 0:
-                return True
-            return False
-
-        self.add_skip_case(
-            teller2,
-            SkipReasons.TRT_NOT_SUPPORT,
-            "INPUT DIM EQUAL TO 1 OF STATIC SHAPE NOT SUPPORT",
-        )
-
-        def teller3(program_config, predictor_config):
             if self.is_int and len(self.dynamic_shape.min_input_shape) == 0:
                 return True
             return False
 
         self.add_skip_case(
-            teller3,
+            teller2,
             SkipReasons.TRT_NOT_SUPPORT,
             "INTEGER INPUT OF STATIC SHAPE NOT SUPPORT",
         )

test/ir/inference/test_trt_convert_softmax.py

@@ -47,8 +47,12 @@ class TrtConvertSoftmaxTest(TrtLayerAutoScanTest):
                 return np.ones([batch, 3, 24]).astype(np.float32)
             elif self.dims == 2:
                 return np.ones([batch, 32]).astype(np.float32)
+            elif self.dims == 1:
+                return np.ones([batch]).astype(np.float32)
+            elif self.dims == 0:
+                return np.ones([]).astype(np.float32)
 
-        for dims in [2, 3, 4]:
+        for dims in [0, 1, 2, 3, 4]:
             for batch in [1, 2, 4]:
                 for axis in [-1, 0, 1, 2, 3]:
                     self.dims = dims
@@ -103,6 +107,14 @@ class TrtConvertSoftmaxTest(TrtLayerAutoScanTest):
                 self.dynamic_shape.min_input_shape = {"softmax_input": [1, 32]}
                 self.dynamic_shape.max_input_shape = {"softmax_input": [4, 64]}
                 self.dynamic_shape.opt_input_shape = {"softmax_input": [1, 32]}
+            elif self.dims == 1:
+                self.dynamic_shape.min_input_shape = {"softmax_input": [1]}
+                self.dynamic_shape.max_input_shape = {"softmax_input": [4]}
+                self.dynamic_shape.opt_input_shape = {"softmax_input": [1]}
+            elif self.dims == 0:
+                self.dynamic_shape.min_input_shape = {"softmax_input": []}
+                self.dynamic_shape.max_input_shape = {"softmax_input": []}
+                self.dynamic_shape.opt_input_shape = {"softmax_input": []}
 
         def clear_dynamic_shape():
             self.dynamic_shape.min_input_shape = {}
@@ -110,6 +122,8 @@ class TrtConvertSoftmaxTest(TrtLayerAutoScanTest):
             self.dynamic_shape.opt_input_shape = {}
 
         def generate_trt_nodes_num(attrs, dynamic_shape):
+            if not dynamic_shape and (self.dims == 1 or self.dims == 0):
+                return 0, 3
             return 1, 2
 
         attrs = [

test/ir/inference/test_trt_convert_square.py

@@ -29,7 +29,9 @@ class TrtConvertSquareTest(TrtLayerAutoScanTest):
 
     def sample_program_configs(self):
         def generate_input1(dims):
-            if dims == 1:
+            if dims == 0:
+                return np.ones([]).astype(np.float32)
+            elif dims == 1:
                 return np.ones([3]).astype(np.float32)
             elif dims == 2:
                 return np.ones([3, 64]).astype(np.float32)
@@ -38,40 +40,42 @@ class TrtConvertSquareTest(TrtLayerAutoScanTest):
             else:
                 return np.ones([1, 3, 64, 64]).astype(np.float32)
 
-        for dims in [1, 2, 3, 4]:
-            for alpha in [1.0, 2.0, 3.0]:
-                self.dims = dims
-
-                ops_config = [
-                    {
-                        "op_type": "square",
-                        "op_inputs": {
-                            "X": ["input_data"],
-                        },
-                        "op_outputs": {"Out": ["output_data"]},
-                        "op_attrs": {},
-                    }
-                ]
-                ops = self.generate_op_config(ops_config)
-
-                program_config = ProgramConfig(
-                    ops=ops,
-                    weights={},
-                    inputs={
-                        "input_data": TensorConfig(
-                            data_gen=partial(generate_input1, dims)
-                        )
+        for dims in [0, 1, 2, 3, 4]:
+            self.dims = dims
+            ops_config = [
+                {
+                    "op_type": "square",
+                    "op_inputs": {
+                        "X": ["input_data"],
                     },
-                    outputs=["output_data"],
-                )
-
-                yield program_config
+                    "op_outputs": {"Out": ["output_data"]},
+                    "op_attrs": {},
+                }
+            ]
+            ops = self.generate_op_config(ops_config)
+
+            program_config = ProgramConfig(
+                ops=ops,
+                weights={},
+                inputs={
+                    "input_data": TensorConfig(
+                        data_gen=partial(generate_input1, dims)
+                    )
+                },
+                outputs=["output_data"],
+            )
+
+            yield program_config
 
     def sample_predictor_configs(
         self, program_config
     ) -> (paddle_infer.Config, List[int], float):
         def generate_dynamic_shape(attrs):
-            if self.dims == 1:
+            if self.dims == 0:
+                self.dynamic_shape.min_input_shape = {"input_data": []}
+                self.dynamic_shape.max_input_shape = {"input_data": []}
+                self.dynamic_shape.opt_input_shape = {"input_data": []}
+            elif self.dims == 1:
                 self.dynamic_shape.min_input_shape = {"input_data": [1]}
                 self.dynamic_shape.max_input_shape = {"input_data": [128]}
                 self.dynamic_shape.opt_input_shape = {"input_data": [64]}
@@ -102,7 +106,7 @@ class TrtConvertSquareTest(TrtLayerAutoScanTest):
             self.dynamic_shape.opt_input_shape = {}
 
         def generate_trt_nodes_num(attrs, dynamic_shape):
-            if not dynamic_shape and self.dims == 1:
+            if not dynamic_shape and (self.dims == 1 or self.dims == 0):
                 return 0, 3
             return 1, 2
 

test/ir/inference/test_trt_convert_stack.py

@@ -32,9 +32,11 @@ class TrtConvertStackTest(TrtLayerAutoScanTest):
         attrs = [
             program_config.ops[i].attrs for i in range(len(program_config.ops))
         ]
-        # The input dimension should be less than the set axis.
+        # axis must be inside [-(rank+1), rank+1)
         if len(inputs['stack_input1'].shape) < attrs[0]['axis']:
             return False
+        if -(len(inputs['stack_input1'].shape) + 1) > attrs[0]['axis']:
+            return False
 
         return True
 
@@ -48,6 +50,8 @@ class TrtConvertStackTest(TrtLayerAutoScanTest):
                 return np.random.random([batch, 24]).astype(np.float32)
             elif self.dims == 1:
                 return np.random.random([24]).astype(np.float32)
+            elif self.dims == 0:
+                return np.random.random([]).astype(np.float32)
 
         def generate_input2(attrs: List[Dict[str, Any]], batch):
             if self.dims == 4:
@@ -58,6 +62,8 @@ class TrtConvertStackTest(TrtLayerAutoScanTest):
                 return np.random.random([batch, 24]).astype(np.float32)
             elif self.dims == 1:
                 return np.random.random([24]).astype(np.float32)
+            elif self.dims == 0:
+                return np.random.random([]).astype(np.float32)
 
         def generate_input3(attrs: List[Dict[str, Any]], batch):
             if self.dims == 4:
@@ -68,8 +74,10 @@ class TrtConvertStackTest(TrtLayerAutoScanTest):
                 return np.random.random([batch, 24]).astype(np.float32)
             elif self.dims == 1:
                 return np.random.random([24]).astype(np.float32)
+            elif self.dims == 0:
+                return np.random.random([]).astype(np.float32)
 
-        for dims in [1, 2, 3, 4]:
+        for dims in [0, 1, 2, 3, 4]:
             for batch in [1, 4]:
                 for axis in [-2, -1, 0, 1, 2, 3]:
                     self.dims = dims
@@ -176,6 +184,22 @@ class TrtConvertStackTest(TrtLayerAutoScanTest):
                     "stack_input2": [24],
                     "stack_input3": [24],
                 }
+            elif self.dims == 0:
+                self.dynamic_shape.min_input_shape = {
+                    "stack_input1": [],
+                    "stack_input2": [],
+                    "stack_input3": [],
+                }
+                self.dynamic_shape.max_input_shape = {
+                    "stack_input1": [],
+                    "stack_input2": [],
+                    "stack_input3": [],
+                }
+                self.dynamic_shape.opt_input_shape = {
+                    "stack_input1": [],
+                    "stack_input2": [],
+                    "stack_input3": [],
+                }
 
         def clear_dynamic_shape():
             self.dynamic_shape.min_input_shape = {}

test/ir/inference/test_trt_convert_sum.py

@@ -37,6 +37,8 @@ class TrtConvertSumTest(TrtLayerAutoScanTest):
                 return np.ones([batch, 24]).astype(np.float32)
             elif self.dims == 1:
                 return np.ones([24]).astype(np.float32)
+            elif self.dims == 0:
+                return np.ones([]).astype(np.float32)
 
         def generate_input2(batch):
             if self.dims == 4:
@@ -47,6 +49,8 @@ class TrtConvertSumTest(TrtLayerAutoScanTest):
                 return np.ones([batch, 24]).astype(np.float32)
             elif self.dims == 1:
                 return np.ones([24]).astype(np.float32)
+            elif self.dims == 0:
+                return np.ones([]).astype(np.float32)
 
         def generate_input3(batch):
             if self.dims == 4:
@@ -57,8 +61,10 @@ class TrtConvertSumTest(TrtLayerAutoScanTest):
                 return np.ones([batch, 24]).astype(np.float32)
             elif self.dims == 1:
                 return np.ones([24]).astype(np.float32)
+            elif self.dims == 0:
+                return np.ones([]).astype(np.float32)
 
-        for dims in [1, 2, 3, 4]:
+        for dims in [0, 1, 2, 3, 4]:
             for batch in [1, 4]:
                 self.dims = dims
                 ops_config = [
@@ -157,6 +163,22 @@ class TrtConvertSumTest(TrtLayerAutoScanTest):
                     "input2": [24],
                     "input3": [24],
                 }
+            elif self.dims == 0:
+                self.dynamic_shape.min_input_shape = {
+                    "input1": [],
+                    "input2": [],
+                    "input3": [],
+                }
+                self.dynamic_shape.max_input_shape = {
+                    "input1": [],
+                    "input2": [],
+                    "input3": [],
+                }
+                self.dynamic_shape.opt_input_shape = {
+                    "input1": [],
+                    "input2": [],
+                    "input3": [],
+                }
 
         def clear_dynamic_shape():
             self.dynamic_shape.min_input_shape = {}
@@ -164,7 +186,7 @@ class TrtConvertSumTest(TrtLayerAutoScanTest):
             self.dynamic_shape.opt_input_shape = {}
 
         def generate_trt_nodes_num(dynamic_shape):
-            if self.dims == 1 and not dynamic_shape:
+            if (self.dims == 1 or self.dims == 0) and not dynamic_shape:
                 return 0, 5
             return 1, 4
 
@@ -205,8 +227,10 @@ class TrtConvertSumTest1(TrtLayerAutoScanTest):
                 return np.ones([batch, 24]).astype(np.float32)
             elif self.dims == 1:
                 return np.ones([24]).astype(np.float32)
+            else:
+                return np.ones([]).astype(np.float32)
 
-        for dims in [1, 2, 3, 4]:
+        for dims in [0, 1, 2, 3, 4]:
             for batch in [1, 4]:
                 self.dims = dims
                 ops_config = [
@@ -263,6 +287,16 @@ class TrtConvertSumTest1(TrtLayerAutoScanTest):
                 self.dynamic_shape.opt_input_shape = {
                     "input1": [24],
                 }
+            elif self.dims == 0:
+                self.dynamic_shape.min_input_shape = {
+                    "input1": [],
+                }
+                self.dynamic_shape.max_input_shape = {
+                    "input1": [],
+                }
+                self.dynamic_shape.opt_input_shape = {
+                    "input1": [],
+                }
 
         def clear_dynamic_shape():
             self.dynamic_shape.min_input_shape = {}
@@ -270,7 +304,7 @@ class TrtConvertSumTest1(TrtLayerAutoScanTest):
             self.dynamic_shape.opt_input_shape = {}
 
         def generate_trt_nodes_num(dynamic_shape):
-            if self.dims == 1 and not dynamic_shape:
+            if (self.dims == 1 or self.dims == 0) and not dynamic_shape:
                 return 0, 3
             return 1, 2
 

test/ir/inference/test_trt_convert_swish.py

@@ -29,7 +29,9 @@ class TrtConvertSwishTest(TrtLayerAutoScanTest):
 
     def sample_program_configs(self):
         def generate_input1(dims, attrs: List[Dict[str, Any]]):
-            if dims == 1:
+            if dims == 0:
+                return np.ones([]).astype(np.float32)
+            elif dims == 1:
                 return np.ones([3]).astype(np.float32)
             elif dims == 2:
                 return np.ones([3, 64]).astype(np.float32)
@@ -38,7 +40,7 @@ class TrtConvertSwishTest(TrtLayerAutoScanTest):
             else:
                 return np.ones([1, 3, 64, 64]).astype(np.float32)
 
-        for dims in [1, 2, 3, 4]:
+        for dims in [0, 1, 2, 3, 4]:
             for beta in [1.0, 2.0, 3.0]:
                 self.dims = dims
 
@@ -73,7 +75,11 @@ class TrtConvertSwishTest(TrtLayerAutoScanTest):
         self, program_config
     ) -> (paddle_infer.Config, List[int], float):
         def generate_dynamic_shape(attrs):
-            if self.dims == 1:
+            if self.dims == 0:
+                self.dynamic_shape.min_input_shape = {"input_data": []}
+                self.dynamic_shape.max_input_shape = {"input_data": []}
+                self.dynamic_shape.opt_input_shape = {"input_data": []}
+            elif self.dims == 1:
                 self.dynamic_shape.min_input_shape = {"input_data": [1]}
                 self.dynamic_shape.max_input_shape = {"input_data": [128]}
                 self.dynamic_shape.opt_input_shape = {"input_data": [64]}
@@ -104,7 +110,7 @@ class TrtConvertSwishTest(TrtLayerAutoScanTest):
             self.dynamic_shape.opt_input_shape = {}
 
         def generate_trt_nodes_num(attrs, dynamic_shape):
-            if self.dims == 1:
+            if (self.dims == 1 or self.dims == 0) and not dynamic_shape:
                 return 0, 3
             return 1, 2