message: Update IPUEstimator validation with cross_replica_mean support
  Summary: FIX T48420
  Reviewers: georgep, #tensorflow, #framework_ip_review_-_any_oss_or_third-party_code_use_has_been_approved, markf
  Maniphest Tasks: T48420
diff: @@ -61,6 +61,7 @@ _INPUT_FN_KEY = "input_fn"
  _BATCH_SIZE_KEY = "batch_size"
  _ASSIGN_ADD_OP = "AssignAddVariableOp"
  _CROSS_REPLICA_SUM_OP = "IpuCrossReplicaSum"
  +_CROSS_REPLICA_MEAN_OP = "IpuCrossReplicaMean"
  _RESOURCE_UPDATE_OP = "ResourceUpdate"
  _HOST_DEVICE = "/device:CPU:0"
  @@ -334,11 +335,13 @@ def _validate_global_step_not_incremented():
  def _validate_replicated_training_graph():
  - def has_cross_replica_sum_op(g):
  - return any(op.type == _CROSS_REPLICA_SUM_OP for op in g.get_operations())
  + def has_cross_replica_reduce_op(g):
  + return any(
  + op.type == _CROSS_REPLICA_SUM_OP or op.type == _CROSS_REPLICA_MEAN_OP
  + for op in g.get_operations())
  graph = ops.get_default_graph()
  - if has_cross_replica_sum_op(graph):
  + if has_cross_replica_reduce_op(graph):
  return
  # Also check inside the resource update `FuncGraph` if there is one.
  @@ -346,12 +349,12 @@ def _validate_replicated_training_graph():
  if op.type == _RESOURCE_UPDATE_OP:
  resource_update_graph = graph._get_function( # pylint: disable=protected-access
  op.get_attr("to_apply").name).graph
  - if has_cross_replica_sum_op(resource_update_graph):
  + if has_cross_replica_reduce_op(resource_update_graph):
  return
  raise ValueError(
  "This is not a valid replicated training graph because no " +
  - _CROSS_REPLICA_SUM_OP +
  + _CROSS_REPLICA_SUM_OP + " or " + _CROSS_REPLICA_MEAN_OP +
  "operations were found. Did you remember to use the " +
  "`tensorflow.python.ipu.optimizers.CrossReplicaOptimizer`?")
  