Handle the so loading in dev mode

tensorflow_addons/register.py

@@ -1,10 +1,9 @@
 import glob
-import os
 from pathlib import Path
 
 import tensorflow as tf
 
-from tensorflow_addons.utils.resource_loader import get_project_root
+from tensorflow_addons.utils.resource_loader import get_path_to_datafile
 
 
 def register_all(keras_objects: bool = True, custom_kernels: bool = True) -> None:
@@ -103,7 +102,7 @@ def register_custom_kernels() -> None:
 
 
 def _get_all_shared_objects():
-    custom_ops_dir = os.path.join(get_project_root(), "custom_ops")
+    custom_ops_dir = get_path_to_datafile("custom_ops", is_so=True)
     all_shared_objects = glob.glob(custom_ops_dir + "/**/*.so", recursive=True)
     all_shared_objects = [x for x in all_shared_objects if Path(x).is_file()]
     return all_shared_objects

tensorflow_addons/utils/resource_loader.py

@@ -31,7 +31,7 @@ def get_project_root():
     return os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
 
 
-def get_path_to_datafile(path):
+def get_path_to_datafile(path, is_so=False):
     """Get the path to the specified file in the data dependencies.
 
     The path is relative to tensorflow_addons/
@@ -42,6 +42,10 @@ def get_path_to_datafile(path):
       The path to the specified data file
     """
     root_dir = get_project_root()
+    if is_so:
+        bazel_bin_dir = os.path.join(os.path.dirname(root_dir), "bazel-bin")
+        if os.path.isdir(bazel_bin_dir):
+            root_dir = os.path.join(bazel_bin_dir, "tensorflow_addons")
     return os.path.join(root_dir, path.replace("/", os.sep))
 
 
@@ -61,7 +65,9 @@ def ops(self):
             )
         if self._ops is None:
             self.display_warning_if_incompatible()
-            self._ops = tf.load_op_library(get_path_to_datafile(self.relative_path))
+            self._ops = tf.load_op_library(
+                get_path_to_datafile(self.relative_path, is_so=True)
+            )
         return self._ops
 
     def display_warning_if_incompatible(self):

tools/run_gpu_tests.sh

@@ -9,4 +9,4 @@ docker build \
        --build-arg TF_VERSION=2.2.0 \
        --build-arg PY_VERSION=3.5 \
        -t tfa_gpu_tests ./
-docker run --rm -t -v cache_bazel:/root/.cache/bazel --gpus=all tfa_gpu_tests
+docker run --rm -t --gpus=all tfa_gpu_tests

tools/testing/build_and_run_tests.sh

@@ -34,5 +34,5 @@ if ! [ -x "$(command -v nvidia-smi)" ]; then
   EXTRA_ARGS="-n auto"
 fi
 
-
+bazel clean
 python -m pytest -v --functions-durations=20 --modules-durations=5 $EXTRA_ARGS ./tensorflow_addons