tf_device.launch() ops. This is because the embedding op is head outside
    compiled and the compile op is wrapped in launch to execute on host
    during TPURewritePass.

    For example, the tf.OpA with the `mini_batch_splits` attribute will be
    moved after _TPUCompileMlir and the first input will use the
    _TPUCompileMlir program output:

    ```mlir

  if (!read_var_op) return;

  // TODO(lyandy): Check if resource (first replica or replicate block arg)
  // shape has not changed in replicate prior to read. Currently after both
  // ResourceOpLiftingPass and TPURewritePass, there should not be any updates
  // to resources prior to their respective ReadVariableOp.
  if (auto block_arg =
          mlir::dyn_cast<BlockArgument>(read_var_op.getResource())) {