namespace {

#define GEN_PASS_DEF_TPUREWRITEPASS
#include "tensorflow/compiler/mlir/tensorflow/transforms/host_runtime/runtime_passes.h.inc"

struct TPURewritePass : public impl::TPURewritePassBase<TPURewritePass> {
  explicit TPURewritePass(llvm::StringRef _module_name)
      : module_name(_module_name) {}

  void runOnOperation() override;

  llvm::StringRef module_name;
};

    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())) {