message: 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
  Maniphest Tasks: T57992
diff: @@ -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
  