From 1f60c85014891303f39a4968a456b92d2ad54077 Mon Sep 17 00:00:00 2001
From: Francois Chollet <fchollet@google.com>
Date: Fri, 23 Oct 2020 16:31:11 -0700
Subject: [PATCH] Make tf2.enabled() available in the __internal__ namespace.
PiperOrigin-RevId: 338765979
Change-Id: Ifa8bbd78f468e0443de38fb4391e831e235e31f3
---
tensorflow/python/tf2.py | 11 +++++++----
.../python/tools/api/generator/api_init_files.bzl | 1 +
.../tools/api/golden/v2/tensorflow.__internal__.pbtxt | 4 ++++
.../api/golden/v2/tensorflow.__internal__.tf2.pbtxt | 7 +++++++
4 files changed, 19 insertions(+), 4 deletions(-)
create mode 100644 tensorflow/tools/api/golden/v2/tensorflow.__internal__.tf2.pbtxt
diff --git a/tensorflow/python/tf2.py b/tensorflow/python/tf2.py
index 4c9d027221f..11d7b6f1aab 100644
--- a/tensorflow/python/tf2.py
+++ b/tensorflow/python/tf2.py
@@ -24,24 +24,27 @@ from __future__ import print_function
import os
+from tensorflow.python.util.tf_export import tf_export
+
+
_force_enable = None
def enable():
- #Enables v2 behaviors.
+ # Enables v2 behaviors.
global _force_enable
_force_enable = True
def disable():
- #Disables v2 behaviors.
+ # Disables v2 behaviors.
global _force_enable
_force_enable = False
+@tf_export("__internal__.tf2.enabled", v1=[])
def enabled():
- #Returns True iff TensorFlow 2.0 behavior should be enabled.
+ # Returns True iff TensorFlow 2.0 behavior should be enabled.
if _force_enable is None:
return os.getenv("TF2_BEHAVIOR", "0") != "0"
-
return _force_enable
diff --git a/tensorflow/python/tools/api/generator/api_init_files.bzl b/tensorflow/python/tools/api/generator/api_init_files.bzl
index eccd39d3f4e..3522cd5d596 100644
--- a/tensorflow/python/tools/api/generator/api_init_files.bzl
+++ b/tensorflow/python/tools/api/generator/api_init_files.bzl
@@ -11,6 +11,7 @@ TENSORFLOW_API_INIT_FILES = [
"__internal__/distribute/multi_process_runner/__init__.py",
"__internal__/test/__init__.py",
"__internal__/test/combinations/__init__.py",
+ "__internal__/tf2/__init__.py",
"__internal__/tracking/__init__.py",
"__operators__/__init__.py",
"audio/__init__.py",
diff --git a/tensorflow/tools/api/golden/v2/tensorflow.__internal__.pbtxt b/tensorflow/tools/api/golden/v2/tensorflow.__internal__.pbtxt
index effcae38787..c71fd575719 100644
--- a/tensorflow/tools/api/golden/v2/tensorflow.__internal__.pbtxt
+++ b/tensorflow/tools/api/golden/v2/tensorflow.__internal__.pbtxt
@@ -16,6 +16,10 @@ tf_module {
name: "test"
mtype: "<type \'module\'>"
}
+ member {
+ name: "tf2"
+ mtype: "<type \'module\'>"
+ }
member {
name: "tracking"
mtype: "<type \'module\'>"
diff --git a/tensorflow/tools/api/golden/v2/tensorflow.__internal__.tf2.pbtxt b/tensorflow/tools/api/golden/v2/tensorflow.__internal__.tf2.pbtxt
new file mode 100644
index 00000000000..6670e45b227
--- /dev/null
+++ b/tensorflow/tools/api/golden/v2/tensorflow.__internal__.tf2.pbtxt
@@ -0,0 +1,7 @@
+path: "tensorflow.__internal__.tf2"
+tf_module {
+ member_method {
+ name: "enabled"
+ argspec: "args=[], varargs=None, keywords=None, defaults=None"
+ }
+}