[XLA] Update dependencies on jit.experimental_jit_scope

tf.contrib is gone in TF 2.0

PiperOrigin-RevId: 265036551

tensorflow/compiler/tests/BUILD

@@ -36,7 +36,6 @@ py_library(
     srcs_version = "PY2AND3",
     visibility = [":friends"],
     deps = [
-        "//tensorflow/contrib/compiler:compiler_py",
         "//tensorflow/core:protos_all_py",
         "//tensorflow/python:array_ops",
         "//tensorflow/python:client",
@@ -46,6 +45,7 @@ py_library(
         "//tensorflow/python:random_seed",
         "//tensorflow/python:session",
         "//tensorflow/python:variables",
+        "//tensorflow/python/compiler/xla:compiler_py",
         "//third_party/py/numpy",
     ],
 )
@@ -1201,7 +1201,7 @@ cuda_py_test(
     srcs = ["jit_test.py"],
     additional_deps = [
         ":test_utils",
-        "//tensorflow/contrib/compiler:compiler_py",
+        "//tensorflow/python/compiler/xla:compiler_py",
         "//tensorflow/core:protos_all_py",
         "//tensorflow/python:array_ops",
         "//tensorflow/python:client",
@@ -1227,7 +1227,7 @@ cuda_py_test(
     srcs = ["dense_layer_test.py"],
     additional_deps = [
         ":test_utils",
-        "//tensorflow/contrib/compiler:compiler_py",
+        "//tensorflow/python/compiler/xla:compiler_py",
         "//tensorflow/core:protos_all_py",
         "//tensorflow/python:array_ops",
         "//tensorflow/python:client_testlib",
@@ -1409,6 +1409,7 @@ tf_xla_py_test(
         ":xla_test",
         "//tensorflow/python:array_ops",
         "//tensorflow/python:framework",
+        "//tensorflow/python:layers",
         "//tensorflow/python:nn",
         "//tensorflow/python:nn_ops",
         "//tensorflow/python:nn_ops_gen",

tensorflow/compiler/tests/dense_layer_test.py

@@ -22,8 +22,8 @@ import os
 import numpy as np
 
 from tensorflow.compiler.tests import test_utils
-from tensorflow.contrib.compiler import jit
 from tensorflow.core.protobuf import config_pb2
+from tensorflow.python.compiler.xla import jit
 from tensorflow.python.layers import layers
 from tensorflow.python.ops import array_ops
 from tensorflow.python.ops import variables

tensorflow/compiler/tests/jit_test.py

@@ -22,10 +22,10 @@ import os
 import numpy as np
 
 from tensorflow.compiler.tests import test_utils
-from tensorflow.contrib.compiler import jit
 from tensorflow.core.protobuf import config_pb2
 from tensorflow.core.protobuf import rewriter_config_pb2
 from tensorflow.python.client import session as session_lib
+from tensorflow.python.compiler.xla import jit
 from tensorflow.python.framework import constant_op
 from tensorflow.python.framework import dtypes
 from tensorflow.python.framework import function

tensorflow/compiler/tests/xla_test.py

@@ -25,12 +25,12 @@ import re
 
 import numpy as np
 
-from tensorflow.python.eager import context
-from tensorflow.contrib.compiler import jit
-from tensorflow.core.protobuf import rewriter_config_pb2
 from tensorflow.core.framework import types_pb2
 from tensorflow.core.protobuf import config_pb2
+from tensorflow.core.protobuf import rewriter_config_pb2
 from tensorflow.python.client import session
+from tensorflow.python.compiler.xla import jit
+from tensorflow.python.eager import context
 from tensorflow.python.framework import dtypes
 from tensorflow.python.framework import ops
 from tensorflow.python.framework import random_seed

tensorflow/python/compiler/xla/BUILD

@@ -79,6 +79,7 @@ cuda_py_test(
         "//tensorflow/python:control_flow_ops",
         "//tensorflow/python:control_flow_util",
         "//tensorflow/python:math_ops",
+        "//tensorflow/python/estimator:estimator_py",
         "//tensorflow/python:platform",
         "//tensorflow/python:state_ops",
         "//tensorflow/python:summary",

tensorflow/python/distribute/BUILD

@@ -697,6 +697,7 @@ tf_xla_py_test(
         ":tpu_strategy",
         "//tensorflow/compiler/tests:xla_test",
         "//tensorflow/python/eager:test",
+        "//tensorflow/python/keras",
         "//tensorflow/python/training/tracking:util",
     ],
 )