tree 8209207da636a0d193f0f8408fae1ffdcdf740a6
parent 6c0a21661fa91eeb4579a55a7a1d511402364e35
author Michael Delorimier <mdel@google.com> 1658949440 -0700
committer TensorFlower Gardener <gardener@tensorflow.org> 1658949631 -0700

`TPUReshardVariables` is inserted even when there is outside compilation. Before this change, all models with a `parallel_execute` were skipped, because `TPUReshardVariables` does not work with model parallelism. Now model parallelism detection is precise: It's detected when there is more than 1 `TPUExecuteAndUpdateVariables`.

PiperOrigin-RevId: 463647960