state, and TPUReshardVariablesOps compares the current state with a desired
state (which can be the compilation result). If they mismatch,
TPUReshardVariablesOp reformats the variables to the desired state; if they
match, TPUReshardVariablesOp is a no-op.

A major use of this pass is weight-update sharding in data parallelism, so we
require there is a tf_device.replicate in the loop.