message: Disable constant folding hook for tf.AllToAll and tf.CrossReplicaSum.
  These TPU ops, under a replicated computation, will have group_assignment > 1 while eager execution/on demand op has access to one TPU core, resulting in a mismatch.
diff: @@ -224,7 +224,7 @@ retained with length 1.
  let verifier = [{ return Verify(*this); }];
  }
  -def TF_AllToAllOp : TF_Op<"AllToAll", [NoSideEffect]> {
  +def TF_AllToAllOp : TF_Op<"AllToAll", [NoSideEffect, TF_NoConstantFold]> {
  let summary = "An Op to exchange data across TPU replicas.";
  let description = [{
  @@ -2154,7 +2154,7 @@ of corresponding 3-element vectors is cross-multiplied independently.
  TF_DerivedOperandTypeAttr T = TF_DerivedOperandTypeAttr<0>;
  }
  -def TF_CrossReplicaSumOp : TF_Op<"CrossReplicaSum", [NoSideEffect, TF_AllTypesMatch<["input", "output"]>]> {
  +def TF_CrossReplicaSumOp : TF_Op<"CrossReplicaSum", [NoSideEffect, TF_AllTypesMatch<["input", "output"]>, TF_NoConstantFold]> {
  let summary = "An Op to sum inputs across replicated TPU instances.";
  let description = [{
  