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

class SimpleTFDeviceAssignmentPass
    : public impl::SimpleTFDeviceAssignmentPassBase<
          SimpleTFDeviceAssignmentPass> {
 public:
  SimpleTFDeviceAssignmentPass() = default;
  SimpleTFDeviceAssignmentPass(const SimpleTFDeviceAssignmentPass&) {}
  explicit SimpleTFDeviceAssignmentPass(llvm::StringRef default_device) {
    default_device_ = std::string(default_device);

  let constructor = "CreateFunctionalToExecutorDialectConversionPass()";
  let dependentDialects = ["mlir::tf_executor::TensorFlowExecutorDialect"];
}

def SimpleTFDeviceAssignmentPass : Pass<"tf-simple-device-assignment", "mlir::func::FuncOp"> {
  let summary = "Simple device assignment in TF dialect.";
  let constructor = "TF::CreateSimpleTFDeviceAssignmentPass()";
  let options = [