From e934a65fd42584f8f624d3e57a9e945b9ad3de46 Mon Sep 17 00:00:00 2001 From: Dan Moldovan Date: Thu, 9 May 2019 16:07:40 -0700 Subject: [PATCH] Fix whitelist bug which incorrectly matched all modules whose name begins with tensorflow, e.g. "tensorflow_foo". Remove tensorboard whitelist. PiperOrigin-RevId: 247514491 --- tensorflow/python/autograph/impl/conversion.py | 8 +------- tensorflow/python/autograph/impl/conversion_test.py | 11 +++++++++++ 2 files changed, 12 insertions(+), 7 deletions(-) diff --git a/tensorflow/python/autograph/impl/conversion.py b/tensorflow/python/autograph/impl/conversion.py index ee472ab5167..6979a1a4847 100644 --- a/tensorflow/python/autograph/impl/conversion.py +++ b/tensorflow/python/autograph/impl/conversion.py @@ -335,16 +335,10 @@ def is_whitelisted_for_graph(o, check_call_override=True): if hasattr(m, '__name__'): # Builtins typically have unnamed modules. for prefix, in config.DEFAULT_UNCOMPILED_MODULES: - if m.__name__.startswith(prefix): + if m.__name__.startswith(prefix + '.') or m.__name__ == prefix: logging.log(2, 'Whitelisted: %s: name starts with "%s"', o, prefix) return True - # Temporary -- whitelist tensorboard modules. - # TODO(b/122731813): Remove. - if m.__name__ == 'tensorboard' or '.tensorboard' in m.__name__: - logging.log(2, 'Whitelisted: %s: name contains "tensorboard"', o) - return True - if hasattr(o, 'autograph_info__') or hasattr(o, '__ag_compiled'): logging.log(2, 'Whitelisted: %s: already converted', o) return True diff --git a/tensorflow/python/autograph/impl/conversion_test.py b/tensorflow/python/autograph/impl/conversion_test.py index bdbdb87e7b0..02efb788e98 100644 --- a/tensorflow/python/autograph/impl/conversion_test.py +++ b/tensorflow/python/autograph/impl/conversion_test.py @@ -18,6 +18,7 @@ from __future__ import absolute_import from __future__ import division from __future__ import print_function +import imp import gast from tensorflow.python.autograph import utils @@ -46,6 +47,16 @@ class ConversionTest(test.TestCase): self.assertTrue(conversion.is_whitelisted_for_graph(utils)) self.assertTrue(conversion.is_whitelisted_for_graph(constant_op.constant)) + def test_is_whitelisted_for_graph_tensorflow_like(self): + + tf_like = imp.new_module('tensorflow_foo') + def test_fn(): + pass + tf_like.test_fn = test_fn + test_fn.__module__ = tf_like + + self.assertFalse(conversion.is_whitelisted_for_graph(tf_like.test_fn)) + def test_convert_entity_to_ast_unsupported_types(self): with self.assertRaises(NotImplementedError): program_ctx = self._simple_program_ctx()