Improved error message: verify that "distributed_batch_norm_replica_group_size" divides replication factor (TF2)

Summary: D73331 cherry-picked for TF2

Reviewers: #tensorflow, #framework_ip_review_-_any_oss_or_third-party_code_use_has_been_approved, gauthamg

Reviewed By: #tensorflow, #framework_ip_review_-_any_oss_or_third-party_code_use_has_been_approved, gauthamg

Maniphest Tasks: T57992

Differential Revision: https://phabricator.sourcevertex.net/D73609

Author: alex-coniasse

tensorflow/compiler/plugin/poplar/driver/poplar_compiler.cc

@@ -1624,6 +1624,8 @@ StatusOr<std::unique_ptr<PoplarExecutableCore>> CompileEngine(
 
   const auto num_local_ipus = poplar_executor->GetNumIpusInLocalProcess(target);
   const auto local_replication_factor = num_local_ipus / num_shards;
+  const auto replica_group_size =
+      poplar_executor->ExperimentalDistributedBatchNormReplicaGroupSize();
 
   if (num_local_ipus % num_shards) {
     return xla::InternalErrorStrCat(
@@ -1632,6 +1634,14 @@ StatusOr<std::unique_ptr<PoplarExecutableCore>> CompileEngine(
         " The number of shards needs to divide the number of local IPUs.");
   }
 
+  if (replica_group_size && replication_factor % replica_group_size) {
+    return xla::InternalErrorStrCat(
+        "The number of replicas (", replication_factor,
+        ") must be divisible by",
+        " distributed_batch_norm_replica_group_size (", replica_group_size,
+        ").");
+  }
+
   // Currently we only support performing replica partitioning across the local
   // replicas in each process, as this allows access to all the parts of a
   // partitioned remote buffer locally. This means that copying to/from all the