}

#define GEN_PASS_DEF_TPUDEVICEPROPAGATIONPASS
#include "tensorflow/compiler/mlir/tensorflow/transforms/tf_passes.h.inc"

struct TPUDevicePropagation
    : public impl::TPUDevicePropagationPassBase<TPUDevicePropagation> {
  void runOnOperation() override;
};

void TPUDevicePropagation::runOnOperation() {
  func::FuncOp func = getOperation();
  if (!IsSupportedGraph(func)) return;