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.

  reformat_operands.push_back(replicate.GetBody().getArgument(
      replicate.GetNumReplicatedBlockArguments() - 1));
  builder.setInsertionPoint(execute_launch);
  auto reformat_op = builder.create<TF::TPUReshardVariablesOp>(
      execute_launch.getLoc(), llvm::ArrayRef<Type>{}, reformat_operands);
  WrapOpInLaunch(&builder, execute_launch.getLoc(), reformat_op,
                 execute_launch.getDevice());

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.