}
        }
      }
    });
  }
  // Remove `kSkipIslandOutlining` attributes.
  for (func::FuncOp func_op : getOperation().getOps<func::FuncOp>()) {
    if (func_op->hasAttr(mlir::TF::kSkipIslandOutlining)) {
      func_op->removeAttr(mlir::TF::kSkipIslandOutlining);
    }
  }
}

}  // namespace

std::unique_ptr<OperationPass<ModuleOp>>

      kTfSavedModelIndexPathAttr, builder.getArrayAttr({file_prefix_attr}))});

  const int insert_idx = func_op.getNumArguments();

  func_op.insertArgument(insert_idx, /*argType=*/filename_op_type, arg_attrs,
                         NameLoc::get(file_prefix_attr));

  return func_op.getArgument(insert_idx);
}

// Creates a 1D string array constant for "tensor_names" input of `RestoreV2`

  // Insert CalibrationStatisticsSaverOp to the end of each region.
  for (auto func_op : module_op.getOps<func::FuncOp>()) {
    int32_t output_file_idx = 0;
    StringRef func_name = func_op.getSymName();

    func_op.walk([&output_file_idx, &ctx, &func_name, &aggregator_ops_to_ignore,
                  this](Operation* op) {
      for (Region& region : op->getRegions()) {

    return failure();
  }

  // Fetch all VarHandleOp.
  llvm::StringSet<> variable_names;
  llvm::SmallVector<TF::VarHandleOp, 4> var_ops;
  for (auto func_op : module.getOps<func::FuncOp>()) {
    for (auto var_handle_op : func_op.getOps<TF::VarHandleOp>()) {
      auto variable_name = GetVariableName(var_handle_op);
      if (variable_names.count(variable_name)) continue;
      var_ops.emplace_back(var_handle_op);

 public:
  void runOnOperation() override;
};

void InputMetricsLoweringPass::runOnOperation() {
  bool has_dynamic_op = false;
  Operation* func_op = getOperation();

  func_op->walk([&](Operation* op) {
    auto abstractOp = op->getRegisteredInfo();
    if (!abstractOp) return WalkResult::advance();

    if (mlir::mhlo::IsDynamicPadderOp(abstractOp->getTypeID())) {

  return op.getAttrOfType<StringAttr>(kXlaOutsideCompilationAttr) != nullptr;
}

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

  auto walk_result = func_op->walk([&](Operation* op) {
    if (IsDeviceClusterOp(*op) && HasChildLaunchDeviceOp(*op)) {
      std::string launch_error =
          absl::StrCat("Node `", op->getName().getStringRef().str(), "` ",

    // Builder builder(func_op.getContext());
    func_op->setAttr(
        kTfInputShapesAttr,
        ArrayAttr::get(func_op.getContext(), updated_input_shapes_attr));
  }
}

// Validates func ops. Returns `failure` if the function is invalid.
LogicalResult ValidateFuncOp(func::FuncOp func_op) {
  auto input_shapes_attr =
      func_op->getAttrOfType<ArrayAttr>(kTfInputShapesAttr);

  ASSERT_TRUE(module_op);

  auto func_op = module_op->lookupSymbol<func::FuncOp>("fully_quantized_add");
  ASSERT_THAT(func_op, NotNull());

  auto add_op_itr = func_op.getBody().op_begin<mlir::stablehlo::AddOp>();
  ASSERT_THAT(add_op_itr,
              Ne(func_op.getBody().op_end<mlir::stablehlo::AddOp>()));

  EXPECT_TRUE(IsOpFullyQuantized(*add_op_itr));
}

    SymbolTable symbol_table(call_op->getParentOfType<ModuleOp>());
    auto func_op =
        dyn_cast_or_null<func::FuncOp>(symbol_table.lookup(func_name));
    if (!func_op) return failure();
    // The quantized fusion should have requantize and return ops at the end.
    auto return_op = dyn_cast_or_null<func::ReturnOp>(
        func_op.getRegion().getBlocks().front().getTerminator());
    if (!return_op) return failure();

  OwningOpRef<ModuleOp> module_op =
      ParseModuleOpString(kArgumentWithUnknownDims);
  ASSERT_TRUE(module_op);

  auto func_op = module_op->lookupSymbol<func::FuncOp>("unknown_dims_arg");
  ASSERT_THAT(func_op, NotNull());
  ASSERT_THAT(func_op.getNumArguments(), Eq(1));

  EXPECT_TRUE(HasRankOf(func_op.getArgument(0), /*rank=*/2));
}

TEST_F(AttrsAndConstraintsTest, HasRankOfReturnsFalseForUnknownRank) {