From 353b8a1adcb471a48ef9b1c5cbfc6097d036473e Mon Sep 17 00:00:00 2001 From: Mihai Maruseac Date: Wed, 23 Oct 2019 10:26:19 -0700 Subject: [PATCH] Add __reduce__ method in virtual pip root due to lazy loading Fixes pickling issue #32159 Tested manually: just applied patch to tensorflow nightly and checked required imports. PiperOrigin-RevId: 276301497 Change-Id: I6b3b6ae8b1218b43c31403ea7cc595ed11136ff9 --- tensorflow/virtual_root_template_v1.__init__.py | 3 +++ tensorflow/virtual_root_template_v2.__init__.py | 3 +++ 2 files changed, 6 insertions(+) diff --git a/tensorflow/virtual_root_template_v1.__init__.py b/tensorflow/virtual_root_template_v1.__init__.py index e3049cee645..236e9f52258 100644 --- a/tensorflow/virtual_root_template_v1.__init__.py +++ b/tensorflow/virtual_root_template_v1.__init__.py @@ -54,6 +54,9 @@ class _LazyLoader(_types.ModuleType): module = self._load() return dir(module) + def __reduce__(self): + return __import__, (self.__name__,) + # Forwarding a module is as simple as lazy loading the module from the new path # and then registering it to sys.modules using the old path diff --git a/tensorflow/virtual_root_template_v2.__init__.py b/tensorflow/virtual_root_template_v2.__init__.py index e09127d70a6..83c020182a8 100644 --- a/tensorflow/virtual_root_template_v2.__init__.py +++ b/tensorflow/virtual_root_template_v2.__init__.py @@ -54,6 +54,9 @@ class _LazyLoader(_types.ModuleType): module = self._load() return dir(module) + def __reduce__(self): + return __import__, (self.__name__,) + # Forwarding a module is as simple as lazy loading the module from the new path # and then registering it to sys.modules using the old path