Update TF OSS API generation script to support OSS keras PIP package.

PiperOrigin-RevId: 359374348
Change-Id: I17b63aea5086250a19b98ad04ad9d17ded3d1f88

tensorflow/api_template.__init__.py

@@ -84,13 +84,22 @@ if _module_dir:
   _current_module.__path__ = [_module_dir] + _current_module.__path__
 setattr(_current_module, "estimator", estimator)
 
-try:
-  from .python.keras.api._v2 import keras
-  _current_module.__path__ = (
-      [_module_util.get_parent_dir(keras)] + _current_module.__path__)
-  setattr(_current_module, "keras", keras)
-except ImportError:
-  pass
+if _os.environ.get("_PREFER_OSS_KERAS", False):
+  try:
+    from keras.api._v2 import keras
+    _current_module.__path__ = (
+        [_module_util.get_parent_dir(keras)] + _current_module.__path__)
+    setattr(_current_module, "keras", keras)
+  except ImportError:
+    pass
+else:
+  try:
+    from .python.keras.api._v2 import keras
+    _current_module.__path__ = (
+        [_module_util.get_parent_dir(keras)] + _current_module.__path__)
+    setattr(_current_module, "keras", keras)
+  except ImportError:
+    pass
 
 # Explicitly import lazy-loaded modules to support autocompletion.
 # pylint: disable=g-import-not-at-top

tensorflow/api_template_v1.__init__.py

@@ -75,13 +75,22 @@ if _module_dir:
   _current_module.__path__ = [_module_dir] + _current_module.__path__
 setattr(_current_module, "estimator", estimator)
 
-try:
-  from .python.keras.api._v1 import keras
-  _current_module.__path__ = (
-      [_module_util.get_parent_dir(keras)] + _current_module.__path__)
-  setattr(_current_module, "keras", keras)
-except ImportError:
-  pass
+if _os.environ.get("_PREFER_OSS_KERAS", False):
+  try:
+    from keras.api._v1 import keras
+    _current_module.__path__ = (
+        [_module_util.get_parent_dir(keras)] + _current_module.__path__)
+    setattr(_current_module, "keras", keras)
+  except ImportError:
+    pass
+else:
+  try:
+    from .python.keras.api._v1 import keras
+    _current_module.__path__ = (
+        [_module_util.get_parent_dir(keras)] + _current_module.__path__)
+    setattr(_current_module, "keras", keras)
+  except ImportError:
+    pass
 
 # Explicitly import lazy-loaded modules to support autocompletion.
 # pylint: disable=g-import-not-at-top

tensorflow/compat_template.__init__.py

@@ -52,13 +52,22 @@ if _module_dir:
   _current_module.__path__ = [_module_dir] + _current_module.__path__
 setattr(_current_module, "estimator", estimator)
 
-try:
-  from tensorflow.python.keras.api._v2 import keras
-  _current_module.__path__ = (
-      [_module_util.get_parent_dir(keras)] + _current_module.__path__)
-  setattr(_current_module, "keras", keras)
-except ImportError:
-  pass
+if _os.environ.get("_PREFER_OSS_KERAS", False):
+  try:
+    from keras.api._v2 import keras
+    _current_module.__path__ = (
+        [_module_util.get_parent_dir(keras)] + _current_module.__path__)
+    setattr(_current_module, "keras", keras)
+  except ImportError:
+    pass
+else:
+  try:
+    from tensorflow.python.keras.api._v2 import keras
+    _current_module.__path__ = (
+        [_module_util.get_parent_dir(keras)] + _current_module.__path__)
+    setattr(_current_module, "keras", keras)
+  except ImportError:
+    pass
 
 # Explicitly import lazy-loaded modules to support autocompletion.
 # pylint: disable=g-import-not-at-top

tensorflow/compat_template_v1.__init__.py

@@ -42,13 +42,22 @@ if _module_dir:
   _current_module.__path__ = [_module_dir] + _current_module.__path__
 setattr(_current_module, "estimator", estimator)
 
-try:
-  from tensorflow.python.keras.api._v1 import keras
-  _current_module.__path__ = (
-      [_module_util.get_parent_dir(keras)] + _current_module.__path__)
-  setattr(_current_module, "keras", keras)
-except ImportError:
-  pass
+if _os.environ.get("_PREFER_OSS_KERAS", False):
+  try:
+    from keras.api._v1 import keras
+    _current_module.__path__ = (
+        [_module_util.get_parent_dir(keras)] + _current_module.__path__)
+    setattr(_current_module, "keras", keras)
+  except ImportError:
+    pass
+else:
+  try:
+    from tensorflow.python.keras.api._v1 import keras
+    _current_module.__path__ = (
+        [_module_util.get_parent_dir(keras)] + _current_module.__path__)
+    setattr(_current_module, "keras", keras)
+  except ImportError:
+    pass
 
 # Explicitly import lazy-loaded modules to support autocompletion.
 # pylint: disable=g-import-not-at-top