/*name=*/kTfQuantInsertRestoreOpStepName,
      /*add_passes_func=*/
      [](mlir::PassManager &pm) {
        pm.addPass(mlir::quant::CreateInsertRestoreOpPass());
        pm.addPass(mlir::quant::CreateInsertSaveOpPass());
        // Initialization by `tf.ConstOp` is no longer required as there is
        // a `tf.RestoreV2Op` now.
        pm.addPass(
            mlir::quant::CreateRemoveVariableInitializationByConstPass());

    return;
  }

  CreateSaveFunc(module_op, target_var_handle_ops);
}

static PassRegistration<InsertSaveOpPass> pass{};

}  // namespace

std::unique_ptr<OperationPass<ModuleOp>> CreateInsertSaveOpPass() {
  return std::make_unique<InsertSaveOpPass>();
}

}  // namespace quant

// `tf.AssignVariableOp(tf.VarHandleOp, tf.Const)` pattern in the initializer
// function of type "restore_op" to identify the VarHandleOps that should be
// saved using the SaveV2 op.
std::unique_ptr<OperationPass<ModuleOp>> CreateInsertSaveOpPass();

// Creates a pass that marks functions with the attribute `tf._noinline = true`
// to avoid being inlined by the `InlinerPass`. `noinline_functions` is the name
// of the functions to mark.