Update all context.graph_context() to ops.get_default_graph().as_default().

PiperOrigin-RevId: 295230258
Change-Id: Id7da36f985d6eae3f9e884e1c8e1352565c4454f

tensorflow/python/keras/engine/base_layer.py

@@ -571,24 +571,22 @@ class Layer(module.Module, version_utils.LayerVersionSelector):
       # with the shape the Layer will be called on (these users will have to
       # implement `compute_output_shape` themselves).
       self._maybe_build(input_shape)
-      with context.graph_mode():
-        graph = func_graph.FuncGraph('graph')
-        with graph.as_default():
-          input_shape = tf_utils.convert_shapes(input_shape, to_tuples=False)
-          def _make_placeholder_like(shape):
-            ph = backend.placeholder(shape=shape, dtype=self.dtype)
-            ph._keras_mask = None
-            return ph
-          inputs = nest.map_structure(_make_placeholder_like, input_shape)
-          try:
-            outputs = self(inputs, training=False)
-          except TypeError as e:
-            six.raise_from(
-                NotImplementedError(
-                    'We could not automatically infer the static shape of the '
-                    'layer\'s output. Please implement the '
-                    '`compute_output_shape` method on your layer (%s).' %
-                    self.__class__.__name__), e)
+      with func_graph.FuncGraph('graph').as_default():
+        input_shape = tf_utils.convert_shapes(input_shape, to_tuples=False)
+        def _make_placeholder_like(shape):
+          ph = backend.placeholder(shape=shape, dtype=self.dtype)
+          ph._keras_mask = None
+          return ph
+        inputs = nest.map_structure(_make_placeholder_like, input_shape)
+        try:
+          outputs = self(inputs, training=False)
+        except TypeError as e:
+          six.raise_from(
+              NotImplementedError(
+                  'We could not automatically infer the static shape of the '
+                  'layer\'s output. Please implement the '
+                  '`compute_output_shape` method on your layer (%s).' %
+                  self.__class__.__name__), e)
       return nest.map_structure(lambda t: t.shape, outputs)
     raise NotImplementedError
 

tensorflow/python/keras/engine/base_layer_test.py

@@ -101,7 +101,7 @@ class BaseLayerTest(keras_parameterized.TestCase):
 
   @keras_parameterized.run_with_all_model_types
   def test_dynamic_layer_error_running_in_graph_mode(self):
-    with context.graph_mode():
+    with ops.get_default_graph().as_default():
       model = testing_utils.get_model_from_layers([DynamicLayer(dynamic=True)],
                                                   input_shape=(3,))
       self.assertEqual(model.dynamic, True)

tensorflow/python/keras/engine/base_layer_v1.py

@@ -535,7 +535,7 @@ class Layer(base_layer.Layer):
       # with the shape the Layer will be called on (these users will have to
       # implement `compute_output_shape` themselves).
       self._maybe_build(input_shape)
-      with context.graph_mode():
+      with ops.get_default_graph().as_default():
         graph = func_graph.FuncGraph('graph')
         with graph.as_default():
           input_shape = tf_utils.convert_shapes(input_shape, to_tuples=False)

tensorflow/python/keras/engine/training_test.py

@@ -154,7 +154,7 @@ class CompileTest(keras_parameterized.TestCase):
     self.assertAllEqual(model._loss_weights_list, [1., 2.])
 
   def test_compile_with_multi_output_and_loss_weights_dict(self):
-    with context.graph_mode():
+    with ops.get_default_graph().as_default():
       model = self._get_multi_output_model()
       loss_weights = {'dense_1': 1., 'dense_2': 2.}
       model.compile(optimizer='adam', loss='mse', loss_weights=loss_weights)
@@ -2142,7 +2142,7 @@ class LossWeightingTest(keras_parameterized.TestCase):
 
   def test_sample_weight_tensor(self):
     """Tests that sample weight may be defined as a tensor in the graph."""
-    with context.graph_mode():
+    with ops.get_default_graph().as_default():
       # Create a simple pass-through model
       input_layer = keras.layers.Input(shape=1, name='input_layer')
       model = keras.Model(inputs=input_layer, outputs=input_layer)

tensorflow/python/keras/keras_parameterized.py

@@ -27,6 +27,7 @@ from absl.testing import parameterized
 from tensorflow.python import keras
 from tensorflow.python import tf2
 from tensorflow.python.eager import context
+from tensorflow.python.framework import ops
 from tensorflow.python.keras import testing_utils
 from tensorflow.python.platform import test
 from tensorflow.python.util import nest
@@ -399,10 +400,11 @@ def run_all_keras_modes(test_or_class=None,
 
 
 def _v1_session_test(f, test_or_class, config, *args, **kwargs):
-  with context.graph_mode(), testing_utils.run_eagerly_scope(False):
-    with testing_utils.experimental_run_tf_function_scope(False):
-      with test_or_class.test_session(use_gpu=True, config=config):
-        f(test_or_class, *args, **kwargs)
+  with ops.get_default_graph().as_default():
+    with testing_utils.run_eagerly_scope(False):
+      with testing_utils.experimental_run_tf_function_scope(False):
+        with test_or_class.test_session(use_gpu=True, config=config):
+          f(test_or_class, *args, **kwargs)
 
 
 def _v2_eager_test(f, test_or_class, *args, **kwargs):

tensorflow/python/keras/layers/lstm_v2_test.py

@@ -1063,8 +1063,9 @@ class LSTMPerformanceTest(test.Benchmark):
                  ' of normal LSTM, got {0:.2f}'.format(v2_vs_normal))
 
   def benchmark_performance_graph(self):
-    with context.graph_mode(), session_lib.Session(config=_config):
-      self._benchmark_performance_with_standard_cudnn_impl()
+    with ops.get_default_graph().as_default():
+      with session_lib.Session(config=_config):
+        self._benchmark_performance_with_standard_cudnn_impl()
 
   def benchmark_performance_eager(self):
     with context.eager_mode():

tensorflow/python/keras/metrics_test.py

@@ -115,7 +115,7 @@ class KerasSumTest(test.TestCase):
     self.assertAlmostEqual(self.evaluate(m.total), 63.75, 2)
 
   def test_sum_graph_with_placeholder(self):
-    with context.graph_mode(), self.cached_session() as sess:
+    with ops.get_default_graph().as_default(), self.cached_session() as sess:
       m = metrics.Sum()
       v = array_ops.placeholder(dtypes.float32)
       w = array_ops.placeholder(dtypes.float32)
@@ -265,7 +265,7 @@ class MeanTest(keras_parameterized.TestCase):
 
   @keras_parameterized.run_all_keras_modes
   def test_mean_graph_with_placeholder(self):
-    with context.graph_mode(), self.cached_session() as sess:
+    with ops.get_default_graph().as_default(), self.cached_session() as sess:
       m = metrics.Mean()
       v = array_ops.placeholder(dtypes.float32)
       w = array_ops.placeholder(dtypes.float32)
@@ -575,7 +575,7 @@ class KerasAccuracyTest(test.TestCase):
     self.assertAlmostEqual(result, 0.93, 2)  # 2.5/2.7
 
   def test_sparse_categorical_accuracy_mismatched_dims_dynamic(self):
-    with context.graph_mode(), self.cached_session() as sess:
+    with ops.get_default_graph().as_default(), self.cached_session() as sess:
       acc_obj = metrics.SparseCategoricalAccuracy(name='my_acc')
       self.evaluate(variables.variables_initializer(acc_obj.variables))
 

tensorflow/python/keras/optimizer_v2/optimizer_v2_test.py

@@ -607,7 +607,7 @@ class OptimizerTest(test.TestCase):
     self.assertLen(var_list(), 4)
 
   def testVarKey(self):
-    with context.graph_mode():
+    with ops.get_default_graph().as_default():
       a = variables.Variable([1., 2.], name='var')
       b = variables.Variable([1.], name='var')
       self.assertTrue(a._in_graph_mode)
@@ -618,7 +618,7 @@ class OptimizerTest(test.TestCase):
       self.assertEqual('var_1', var_key)
 
   def testVarName(self):
-    with context.graph_mode():
+    with ops.get_default_graph().as_default():
       var = variables.Variable([1., 2.], name='var')
       loss = var + 1.
       opt = adam.Adam()

tensorflow/python/keras/saving/hdf5_format_test.py

@@ -1007,7 +1007,7 @@ class TestWeightSavingAndLoadingTFFormat(test.TestCase):
       model.load_weights(fname)
 
   def test_no_graph_pollution(self):
-    with context.graph_mode():
+    with ops.get_default_graph().as_default():
       graph = ops.Graph()
       with graph.as_default(), self.session(graph) as session:
         model = SubclassedModel()

tensorflow/python/keras/utils/version_utils_test.py

@@ -24,7 +24,6 @@ import numpy as np
 import six
 
 from tensorflow.python import keras
-from tensorflow.python.eager import context
 from tensorflow.python.framework import ops
 from tensorflow.python.keras import keras_parameterized
 from tensorflow.python.keras.engine import base_layer
@@ -152,7 +151,7 @@ class SplitUtilsTest(keras_parameterized.TestCase):
     model = keras.Sequential([keras.layers.Dense(1)])
     model.compile('sgd', 'mse')
     x, y = np.ones((10, 10)), np.ones((10, 1))
-    with context.graph_mode():
+    with ops.get_default_graph().as_default():
       with self.assertRaisesRegexp(
           ValueError, 'instance was constructed with eager mode enabled'):
         model.fit(x, y, batch_size=2)