#include "tensorflow/core/public/session.h"

namespace mlir {
namespace tf_saved_model {

// Initializes all variables in Session Init function for all variables in
// 'session'.
LogicalResult InitializeVariablesInSessionInitializer(
    ModuleOp module, tensorflow::Session *session);

}  // namespace tf_saved_model

}  // namespace mlir

          InitializeVariablesInSessionInitializerPass> {
 public:
  void runOnOperation() final {
    static tensorflow::Session* session = new TF::test_util::FakeSession();
    if (failed(tf_saved_model::InitializeVariablesInSessionInitializer(
            getOperation(), session)))
      signalPassFailure();
  }

 private:
};
}  // anonymous namespace

std::unique_ptr<OperationPass<ModuleOp>>

            .getValue());
  } else {
    return CreateSessionInitFunc(module);
  }
}

}  // namespace

LogicalResult InitializeVariablesInSessionInitializer(
    ModuleOp module, tensorflow::Session* session) {
  const tensorflow::DeviceMgr* mgr = nullptr;
  auto status = session->LocalDeviceManager(&mgr);
  if (!status.ok()) {
    module->emitError(

    // imported back to MLIR. When importing the graph to MLIR, the name of the
    // op is used to retrieve the tensor values of each variable. See
    // `InitializeVariablesInSessionInitializer` for further details.
    const auto loc = NameLoc::get(StringAttr::get(
        var_handle_op.getContext(), var_handle_op.getSharedName()));
    var_handle_op->setLoc(loc);

      return diag_handler.Combine(
          errors::Internal("Failed to lift variables."));
  } else {
    if (failed(mlir::tf_saved_model::InitializeVariablesInSessionInitializer(
            module, bundle.GetSession())))
      return diag_handler.Combine(
          errors::Internal("Failed to initialize variables in session init."));
  }

  pm.clear();