[Zero-Dim] Make auto parallel judge dim more strict (#47961)

paddle/fluid/operators/batch_norm_op.cc

@@ -164,7 +164,7 @@ void BatchNormOp::InferShape(framework::InferShapeContext *ctx) const {
   ctx->SetOutputDim("SavedMean", {C});
   ctx->SetOutputDim("SavedVariance", {C});
   ctx->ShareLoD("X", "Y");
-  if (ctx->HasInput("ReserveSpace")) {
+  if (ctx->HasOutput("ReserveSpace")) {
     ctx->SetOutputDim("ReserveSpace", {-1});
   }
 }

python/paddle/distributed/auto_parallel/completion.py

@@ -1239,7 +1239,7 @@ class Completer:
                                 input_var
                             ).dims_mapping
                     else:
-                        if fwd_op_dist_attr.get_input_dims_mapping(input_name):
+                        if input_name in forward_op.input_arg_names:
                             ref_dims_mapping = (
                                 fwd_op_dist_attr.get_input_dims_mapping(
                                     input_name
@@ -1544,7 +1544,7 @@ class Completer:
                                 input_var
                             ).dims_mapping
                     else:
-                        if fwd_op_dist_attr.get_input_dims_mapping(input_name):
+                        if input_name in forward_op.input_arg_names:
                             ref_dims_mapping = (
                                 fwd_op_dist_attr.get_input_dims_mapping(
                                     input_name