{producer, min_consumer, bad_consumers})));
    break;
  }

  if (failed(applyFullConversion(module.getOperation(), target, finalPatterns)))
    signalPassFailure();

  // The uses of arg control dependency has been dropped. We can safely remove
  // the block argument here.
  module.walk([&](tf_executor::GraphOp graph) {
    graph.getBody().front().eraseArguments(

  addPatterns(context, stage1Patterns, this->preserve_assert_op_);

  FrozenRewritePatternSet stage1FrozenPatterns(std::move(stage1Patterns));
  if (!applyPatterns(func, target, stage1FrozenPatterns)) {
    return signalPassFailure();
  }
  // Explict BroadcastTo addition for left-over broadcast-able ops.
  // The following pattern matchings should be done after the other legalization
  // rules in order not to add unnecessary BroadcastTo ops.

void LegalizeTFCommunication::runOnOperation() {
  auto module = getOperation();
  llvm::SmallDenseMap<StringRef, FuncToRewrite> funcs_to_rewrite;
  if (failed(GetFunctionsToRewrite(module, funcs_to_rewrite)))
    return signalPassFailure();

  OpBuilder builder(&getContext());
  for (const auto& func_and_name : funcs_to_rewrite) {
    const auto& func_to_rewrite = func_and_name.getSecond();
    func::FuncOp func = func_to_rewrite.original;

    if (IsTpuRegularOp(op)) {
      success &= IsValidMAXIMALSharding(op, metadata_map);
      success &= IsValidShardingTupleForArity(op);
    }
    success &= !HasSingleCoreTpu(op);
    if (!success) {
      signalPassFailure();
    }
  });
}

}  // anonymous namespace

std::unique_ptr<OperationPass<ModuleOp>> CreateTPUValidateInputsPass() {
  return std::make_unique<TPUValidateInputsPass>();
}

  RuntimeDevices devices;
  if (failed(::tensorflow::GetDevicesFromOp(func->getParentOfType<ModuleOp>(),
                                            &devices)))
    return signalPassFailure();

  // If there is no runtime device information and data format is not explicitly
  // forced, there is nothing to do.
  if (devices.NumDevices() == 0 && force_data_format_.empty()) return;

  if (!shouldKeepUnusedQdqPattern()) {
    RewritePatternSet patterns_2(&getContext());
    patterns_2.add<RemoveUnusedQdqPattern>(ctx);
    if (failed(applyPatternsAndFoldGreedily(func, std::move(patterns_2)))) {
      signalPassFailure();
    }
  }
}
}  // namespace

// Creates an instance of the TensorFlow dialect Quantize pass.
std::unique_ptr<OperationPass<func::FuncOp>> CreateQuantizePass() {
  QuantizationSpecs quant_specs;

              execute_launch, resource_analysis_info,
              /*check_device=*/!parent_is_replicate,
              /*check_same_region=*/parent_is_replicate, &builder))) {
        signalPassFailure();
        return;
      }
    }
  });
}

}  // namespace

std::unique_ptr<OperationPass<ModuleOp>>
CreateTPUMergeVariablesWithExecutePass() {

  auto result = module.walk([&](tf_device::ClusterOp cluster) {
    OptimizeIfRegions(&cluster.GetBody(), module);
    return WalkResult::advance();
  });
  if (result.wasInterrupted()) return signalPassFailure();
}

}  // namespace

std::unique_ptr<OperationPass<ModuleOp>> CreateMergeControlFlowPass() {
  return std::make_unique<MergeControlFlowPass>();
}

}  // namespace TFDevice

  auto module = getOperation();
  auto main = module.lookupSymbol<func::FuncOp>("main");
  if (!main) return;
  if (failed(DecomposeStackOps(&main.front(), module))) {
    signalPassFailure();
  }
}

}  // namespace

namespace TF {
std::unique_ptr<OperationPass<ModuleOp>> CreateStackOpsDecompositionPass() {
  return std::make_unique<StackOpsDecompositionPass>();
}

    patterns.add(MergeConsecutivePad);
    if (failed(applyPatternsAndFoldGreedily(getOperation(),
                                            std::move(patterns)))) {
      return signalPassFailure();
    }
  }
};

std::unique_ptr<Pass> createOptimizePass() {
  return std::make_unique<OptimizePass>();
}

static PassRegistration<OptimizePass> pass;

}  // namespace odml