PaddlePaddle
diff --git a/‎paddle/phi/api/yaml/backward.yaml
Copy file name to clipboard
+1Lines changed: 1 addition & 0 deletions b/‎paddle/phi/api/yaml/backward.yaml
Copy file name to clipboard
+1Lines changed: 1 addition & 0 deletions
diff --git a/‎paddle/phi/api/yaml/ops.yaml
Copy file name to clipboard
+2Lines changed: 2 additions & 0 deletions b/‎paddle/phi/api/yaml/ops.yaml
Copy file name to clipboard
+2Lines changed: 2 additions & 0 deletions
diff --git a/‎paddle/phi/infermeta/spmd_rules/numel.cc
Copy file name to clipboard
+54Lines changed: 54 additions & 0 deletions b/‎paddle/phi/infermeta/spmd_rules/numel.cc
Copy file name to clipboard
+54Lines changed: 54 additions & 0 deletions
diff --git a/‎paddle/phi/infermeta/spmd_rules/numel.h
Copy file name to clipboard
+32Lines changed: 32 additions & 0 deletions b/‎paddle/phi/infermeta/spmd_rules/numel.h
Copy file name to clipboard
+32Lines changed: 32 additions & 0 deletions
diff --git a/‎paddle/phi/infermeta/spmd_rules/rules.h
Copy file name to clipboard
+6Lines changed: 6 additions & 0 deletions b/‎paddle/phi/infermeta/spmd_rules/rules.h
Copy file name to clipboard
+6Lines changed: 6 additions & 0 deletions
@@ -2602,6 +2602,7 @@
   infer_meta :
     func : GeneralBinaryGradInferMeta
     param : [x, y]
+    spmd_rule: WhereGradInferSpmd
   kernel :
     func : where_grad
   no_need_buffer : x, y
 
@@ -1914,6 +1914,7 @@
   output : Tensor(size)
   infer_meta :
     func : NumelInferMeta
+    spmd_rule : NumelInferSpmd
   kernel :
     func : numel
     data_type : x
@@ -2846,6 +2847,7 @@
   output : Tensor(out)
   infer_meta :
     func : WhereInferMeta
+    spmd_rule: WhereInferSpmd
   kernel :
     func : where
   inplace: (x -> out)
 
@@ -0,0 +1,54 @@
+/* Copyright (c) 2023 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
+
+    https://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. */
+
+#include "paddle/phi/infermeta/spmd_rules/numel.h"
+
+#include "glog/logging.h"
+
+#include "paddle/phi/core/enforce.h"
+#include "paddle/phi/infermeta/spmd_rules/utils.h"
+
+namespace phi {
+namespace distributed {
+using phi::distributed::auto_parallel::str_join;
+
+SpmdInfo NumelInferSpmd(const DistMetaTensor& x) {
+  std::string alphabet = "abcdefghijklmnopqrstuvwxyz";
+  auto x_shape = phi::vectorize(x.dims());
+  int x_ndim = x_shape.size();
+  auto x_dist_attr_src = x.dist_attr();
+  std::vector<int64_t> x_dims_mapping = x_dist_attr_src.dims_mapping();
+  PADDLE_ENFORCE_EQ(
+      x_ndim,
+      x_dims_mapping.size(),
+      phi::errors::InvalidArgument("The Tensor Input's rank [%d] and Input's "
+                                   "dims_mapping size [%d] are not matched.",
+                                   x_ndim,
+                                   x_dims_mapping.size()));
+  TensorDistAttr out_dist_attr;
+  out_dist_attr.set_dims_mapping({});
+  std::vector<int64_t> partial_on_dims;
+  const auto& dim_mapping = x_dims_mapping;
+  for (int i = 0; i < x_ndim; ++i) {
+    auto mapping = dim_mapping[i];
+    if (mapping != -1) {
+      partial_on_dims.push_back(mapping);
+    }
+  }
+  out_dist_attr.set_partial_status(partial_on_dims);
+  return SpmdInfo({x_dist_attr_src}, {out_dist_attr});
+}
+
+}  // namespace distributed
+}  // namespace phi
@@ -0,0 +1,32 @@
+/* Copyright (c) 2023 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
+
+https://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 <iterator>
+#include <map>
+#include <string>
+#include <vector>
+
+#include "paddle/phi/common/int_array.h"
+#include "paddle/phi/core/distributed/auto_parallel/dist_meta_tensor.h"
+#include "paddle/phi/core/distributed/type_defs.h"
+
+namespace phi {
+namespace distributed {
+
+SpmdInfo NumelInferSpmd(const DistMetaTensor& x);
+
+}  // namespace distributed
+}  // namespace phi
@@ -23,6 +23,7 @@ limitations under the License. */
 #include "paddle/phi/infermeta/spmd_rules/flatten.h"
 #include "paddle/phi/infermeta/spmd_rules/layer_norm.h"
 #include "paddle/phi/infermeta/spmd_rules/matmul.h"
+#include "paddle/phi/infermeta/spmd_rules/numel.h"
 #include "paddle/phi/infermeta/spmd_rules/reduction.h"
 #include "paddle/phi/infermeta/spmd_rules/replicated.h"
 #include "paddle/phi/infermeta/spmd_rules/reshape.h"
@@ -32,6 +33,7 @@ limitations under the License. */
 #include "paddle/phi/infermeta/spmd_rules/stack.h"
 #include "paddle/phi/infermeta/spmd_rules/transpose.h"
 #include "paddle/phi/infermeta/spmd_rules/unsqueeze.h"
+#include "paddle/phi/infermeta/spmd_rules/where.h"
 
 /**
  * Design Notes:
@@ -555,5 +557,9 @@ PD_REGISTER_SPMD_RULE(log_softmax,
                       PD_INFER_SPMD(phi::distributed::SoftmaxInferSpmd),
                       PD_INFER_SPMD(phi::distributed::SoftmaxInferSpmdReverse));
 
+PD_REGISTER_SPMD_RULE(where,
+                      PD_INFER_SPMD(phi::distributed::WhereInferSpmd),
+                      PD_INFER_SPMD(phi::distributed::WhereInferSpmdReverse));
+
 }  // namespace distributed
 }  // namespace phi