module.walk([&](tf_device::ClusterFuncOp cluster_func_op) {
    RewriteCall(cluster_func_op, symtab, builder);
  });
}

}  // namespace

namespace TFDevice {
std::unique_ptr<OperationPass<ModuleOp>> CreateXlaRewritePass() {
  return std::make_unique<XlaRewritePass>();
}

}  // namespace TFDevice

  }
}

void AddNonTPULowerClusterToRuntimeOpsPassPipeline(
    OpPassManager& pm, llvm::StringRef module_name) {
  // Rewrite cluster functions into XLA launch ops.
  pm.addPass(mlir::TFDevice::CreateXlaRewritePass());
  // Re-run the canonicalizer pass as some cleanup during resource op lifting
  // pass opens up some opportunities for canonicalization of cluster ops.

      return %arg0 : tensor<i32>
    }
    ```
    Notice that the called function is rewritten, with the order of its parameters changed.
  }];

  let constructor = "TFDevice::CreateXlaRewritePass()";
  let dependentDialects = ["tf_device::TensorFlowDeviceDialect"];
}

def XlaValidateInputsPass : Pass<"tf-xla-validate-inputs", "ModuleOp"> {
  let summary = "Validtes inputs to the TF CPU/GPU bridge";

std::unique_ptr<OperationPass<ModuleOp>> CreateXlaInlineDeviceOpsPass();

// Creates a pass that rewrites partitioned calls with `_xla_compile_device
// type` with `tf.XlaLaunch` ops.
std::unique_ptr<OperationPass<ModuleOp>> CreateXlaRewritePass();

// Create a pass that validates the input graph to the CPU/GPU bridge.
std::unique_ptr<OperationPass<ModuleOp>> CreateXlaValidateInputsPass();
}  // namespace TFDevice

namespace TFTPU {