public:
  void runOnOperation() override;
};

void InferenceMetricsPass::runOnOperation() {
  bool has_tpu_partitioned_call = false;
  ModuleOp module = getOperation();

  for (auto func_op : module.getOps<func::FuncOp>()) {
    func_op->walk(
        [&](TF::TPUPartitionedCallOp op) { has_tpu_partitioned_call = true; });

    if (has_tpu_partitioned_call) break;
  }

  std::string has_tpu_partitioned_call_str =

#include "llvm/ADT/StringRef.h"
#include "mlir/Dialect/Func/IR/FuncOps.h"  // from @llvm-project

namespace mlir {
namespace quant {

// Adds a new input name to the `inputs` field of the `tf.entry_function`
// attribute if the attribute exist in the given function. Otherwise, no
// attribute is modified.
void AddEntryFunctionInput(StringRef input_name, func::FuncOp func_op);

}  // namespace quant
}  // namespace mlir

    EXPECT_TRUE(module_op_ref);
    return module_op_ref;
  }

  mlir::MLIRContext ctx_{};
};

func::FuncOp GetMainFuncOp(ModuleOp module_op) {
  for (auto func_op : module_op.getOps<func::FuncOp>()) {
    if (func_op.getSymName() == "main") {
      return func_op;
    }
  }
  return {};
}

TEST_F(ConvertAssetArgsTest, ConvertsSingleAssetArg) {

// function.
bool ShouldIncludeInMainFunction(func::FuncOp func_op) {
  return !func_op.isPrivate() && IsExported(func_op) &&
         IsEntryFunction(func_op);
}

// Sets a function to be private so it can be referred internally.
void SetFunctionPrivate(func::FuncOp func) {
  func.setVisibility(SymbolTable::Visibility::Private);

// Returns a list of "tf.entry_function" attribute's "inputs" comma-split
// values.
//
// Example: if `func_op` has attribute `tf.entry_function = {inputs =
// "arg0:0,arg1:0"}`, then this function returns `{"arg0:0", "arg1:0"}`.
SmallVector<StringRef> GetEntryFunctionInputs(func::FuncOp func_op) {
  auto entry_function_attr =
      func_op->getAttrOfType<DictionaryAttr>("tf.entry_function");

  SmallVector<StringRef> inputs;

  }

  Value func_input_arg;
  // Creates a function.
  func::FuncOp func_op = PrepareFunctionRegister(
      rewriter, input_val, result_type, unique_func_name, func_input_arg);

  // Fills the body.
  Operation* last_op_in_func =
      quantization_operations_func(rewriter, func_op.getOperation(),
                                   func_input_arg, result_type, quant_type);

    ModuleOp module_op = getOperation();
    SmallVector<func::FuncOp, 4> original_func;
    for (auto func_op : module_op.getOps<func::FuncOp>()) {
      original_func.push_back(func_op);
    }
    for (auto func_op : original_func) {
      std::optional<llvm::StringMap<FunctionMetadata>> metadatas =
          GetFunctionMetadatas(func_op);
      if (!metadatas) {
        signalPassFailure();
        return;

    : public impl::VerifyClusteringPassBase<VerifyClusteringPass> {
 public:
  void runOnOperation() override;
};

void VerifyClusteringPass::runOnOperation() {
  Operation* func_op = getOperation();

  auto walk_result = func_op->walk([&](Operation* op) {
    if (!tensorflow::tf2xla::internal::IsInBridgeAcceptableDialects(op)) {
      std::string error = "op is in dialect " +

  get_alg_op->erase();
  return success();
}

void TFXLADeviceSpecificTransforms::runOnOperation() {
  if (!device_type_.hasValue()) return;
  auto func_op = getOperation();

  auto walk_result = func_op->walk([&](TF::StatelessRandomGetAlgOp op) {
    if (failed(ConvertGetAlgOp(op))) {
      op->emitOpError(
          "Could not convert and remove Device specific information");

  SmallVector<int64_t, 1> layer_norm_slice_shape_;
  SmallVector<int64_t, 1> layer_norm_size_values_;
};

LogicalResult ConvertKerasLSTMLayer(mlir::func::FuncOp func_op,
                                    OpBuilder* builder);

LogicalResult ConvertKerasLSTMLayer(mlir::func::FuncOp func_op,
                                    OpBuilder* builder, bool indy);

}  // end namespace TFL
}  // end namespace mlir