#define GEN_PASS_DEF_REMOVEVARIABLESINSESSIONINITIALIZERPASS
#include "tensorflow/compiler/mlir/tensorflow/transforms/tf_savedmodel_passes.h.inc"

class RemoveVariablesInSessionInitializerPass
    : public impl::RemoveVariablesInSessionInitializerPassBase<
          RemoveVariablesInSessionInitializerPass> {
 public:
  void runOnOperation() override;
};

void RecursiveRemove(Operation* op,

  // `lift_variables == true`) or newly adding variable initialization patterns
  // in the initializer functions. One might want to set this to `true` because
  // the `RemoveVariablesInSessionInitializerPass` pass, which runs otherwise,
  // may unexpectedly also remove the initialization patterns for non-variable
  // resources (like hash tables) if they involve variables. Such a case is
  // illustrated in the test file

  let summary = "Optimize tf_saved_model.global_tensor's.";

  let constructor = "mlir::tf_saved_model::CreateOptimizeGlobalTensorsPass()";
}

def RemoveVariablesInSessionInitializerPass : Pass<"tf-saved-model-remove-vars-in-session-initializer", "ModuleOp"> {
  let summary = "Remove variables in tf saved model's session initializer.";

  // If `include_variables_in_initializers` is set to false, then it removes all
  // variables from the initializer functions (registered in the
  // `tf_saved_model::SessionInitializerOp`) by running the
  // `RemoveVariablesInSessionInitializerPass`, regardless of whether
  // `lift_variable_ops_to_args` is true or not.
  static Status LiftVariables(const SavedModelBundle& bundle,
                              mlir::ModuleOp module,