struct LegalizeWhilePass
    : public impl::LegalizeWhilePassBase<LegalizeWhilePass> {
  MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(LegalizeWhilePass)
  void RunOnFunction(func::FuncOp func);

  void runOnOperation() override {
    for (auto op : getOperation().getOps<func::FuncOp>()) RunOnFunction(op);
  }
};

}  // namespace

// Inserts call to the given function into the 'region'.

// aggregate analysis and operate on all non-extern functions (similar to a
// OperationPass<func::FuncOp>, but with no concurrency between functions). The
// derived classes need to provide a runOnFunction() method that accepts the
// function and the analysis information for that function.
template <typename DerivedT, typename AnalysisT>
class PerFunctionAggregateAnalysisConsumerPass
    : public PassWrapper<

  using TacPass<func::FuncOp>::TacPass;

  ~TacFunctionPass() override = default;

  mlir::func::FuncOp getFunction() { return getOperation(); }

  virtual void runOnFunction() = 0;

  void runOnOperation() final {
    if (!getFunction().isExternal()) runOnFunction();
  }

 protected:
  // Returns the derived pass name.
  StringRef getName() const override { return llvm::getTypeName<T>(); }

      : TacFunctionPass(nullptr) {
    device_specs_flag_ = device_specs;
  }

  explicit TargetAnnotationPass(const TacModule* module)
      : TacFunctionPass(module) {}

 private:
  void runOnFunction() override;
  void SetTargetAnnotation(Operation* op,
                           llvm::ArrayRef<std::string> device_specs,
                           OpBuilder* builder);

  ListOption<std::string> device_specs_flag_{

    : public TF::PerFunctionAggregateAnalysisConsumerPass<
          TestSideEffectAnalysisPass, TF::SideEffectAnalysis> {
  MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(TestSideEffectAnalysisPass)

  void runOnFunction(func::FuncOp func,
                     const TF::SideEffectAnalysis::Info& analysis) {
    int64_t next_id = 0;
    llvm::SmallDenseMap<Operation*, int64_t, 8> ids;
    func.walk([&](Operation* op) {

    return "tf-test-resource-alias-analysis";
  }

  StringRef getDescription() const final {
    return "Add remarks based on resource alias analysis result, for testing "
           "purpose.";
  }

  void runOnFunction(func::FuncOp func,
                     const TF::ResourceAliasAnalysis::Info& analysis) {
    int64_t next_id = 0;
    llvm::SmallDenseMap<Value, int64_t, 8> ids;

    auto assign_id = [&](Value value) {

struct TPUDynamicLayoutPass
    : public TF::PerFunctionAggregateAnalysisConsumerPass<
          TPUDynamicLayoutPass, TF::ResourceAliasAnalysis> {
  MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(TPUDynamicLayoutPass)

  void runOnFunction(
      func::FuncOp func,
      const TF::ResourceAliasAnalysis::Info& resource_alias_analysis);

  StringRef getArgument() const final { return "tf-tpu-dynamic-layout-pass"; }

  StringRef getArgument() const final { return "tf-executor-break-up-islands"; }

  StringRef getDescription() const final {
    return "Transform from TF control dialect to TF executor dialect.";
  }

  void runOnFunction(func::FuncOp func,
                     const TF::SideEffectAnalysis::Info& side_effect_analysis);

  void BreakUpIsland(tf_executor::IslandOp island_op,