1 files changed, 20 insertions, 4 deletions
diff --git a/tensorflow/tensorflow.bzl b/tensorflow/tensorflow.bzl
index 9e26b4c654..c3deecdff6 100644
--- a/tensorflow/tensorflow.bzl
+++ b/tensorflow/tensorflow.bzl
@@ -539,6 +539,7 @@ def _py_wrap_cc_impl(ctx):
   for dep in ctx.attr.deps:
     inputs += dep.cc.transitive_headers
   inputs += ctx.files._swiglib
+  inputs += ctx.files.toolchain_deps
   swig_include_dirs = set(_get_repository_roots(ctx, inputs))
   swig_include_dirs += sorted([f.dirname for f in ctx.files._swiglib])
   args = ["-c++",
@@ -573,6 +574,9 @@ _py_wrap_cc = rule(
             allow_files = True,
             providers = ["cc"],
         ),
+        "toolchain_deps": attr.label_list(
+            allow_files = True,
+        ),
         "module_name": attr.string(mandatory = True),
         "py_module_name": attr.string(mandatory = True),
         "_swig": attr.label(
@@ -763,6 +767,7 @@ def tf_py_wrap_cc(name, srcs, swig_includes=[], deps=[], copts=[], **kwargs):
               srcs=srcs,
               swig_includes=swig_includes,
               deps=deps + extra_deps,
+              toolchain_deps=["//tools/defaults:crosstool"],
               module_name=module_name,
               py_module_name=name)
   extra_linkopts = select({
@@ -812,6 +817,14 @@ def tf_py_wrap_cc(name, srcs, swig_includes=[], deps=[], copts=[], **kwargs):
                       "//conditions:default": [":" + cc_library_name],
                     }))
 
+def py_test(deps=[], **kwargs):
+  native.py_test(
+      deps=select({
+          "//conditions:default" : deps,
+          "//tensorflow:no_tensorflow_py_deps" : []
+      }),
+      **kwargs)
+
 def tf_py_test(name, srcs, size="medium", data=[], main=None, args=[],
                tags=[], shard_count=1, additional_deps=[], flaky=0):
   native.py_test(
@@ -824,10 +837,13 @@ def tf_py_test(name, srcs, size="medium", data=[], main=None, args=[],
       visibility=["//tensorflow:internal"],
       shard_count=shard_count,
       data=data,
-      deps=[
-          "//tensorflow/python:extra_py_tests_deps",
-          "//tensorflow/python:gradient_checker",
-      ] + additional_deps,
+      deps=select({
+          "//conditions:default" : [
+            "//tensorflow/python:extra_py_tests_deps",
+            "//tensorflow/python:gradient_checker",
+          ] + additional_deps,
+          "//tensorflow:no_tensorflow_py_deps" : []
+      }),
       flaky=flaky,
       srcs_version="PY2AND3")