Set DCMAKE_PREFIX_PATH for TensorFlow build to build with pip pybind11 (#143)

* Fix pybind11 install doc
Signed-off-by: Trevor Morris <tmorris@nvidia.com>

* Set CMAKE_PREFIX_PATH for TF to find pybind11
Signed-off-by: Trevor Morris <tmorris@nvidia.com>

* Update test builds to use pip install of apt.
Signed-off-by: Trevor Morris <tmorris@nvidia.com>

---------
Signed-off-by: Trevor Morris <tmorris@nvidia.com>

.github/workflows/build.yml

@@ -43,8 +43,7 @@ jobs:
         uses: actions/checkout@v3
       - name: 'Build'
         run: |
-          apt-get -y update && \
-          apt -y install ninja-build pybind11-dev && \
+          pip install ninja pybind11 && \
           pip install --upgrade "jax[cuda12_local]" -f https://storage.googleapis.com/jax-releases/jax_cuda_releases.html && \
           mkdir -p wheelhouse && \
           NVTE_FRAMEWORK=jax pip wheel -w wheelhouse . -v
@@ -69,8 +68,7 @@ jobs:
         uses: actions/checkout@v3
       - name: 'Build'
         run: |
-          apt-get -y update && \
-          apt -y install ninja-build pybind11-dev && \
+          pip install ninja pybind11 && \
           mkdir -p wheelhouse && \
           NVTE_FRAMEWORK=tensorflow pip wheel -w wheelhouse . -v
       - name: 'Upload wheel'

setup.py

@@ -161,7 +161,8 @@ class PyTorchBuilder(FrameworkBuilderBase):
 
 class TensorFlowBuilder(FrameworkBuilderBase):
     def cmake_flags(self):
-        return ["-DENABLE_TENSORFLOW=ON"]
+        p = [d for d in sys.path if 'dist-packages' in d][0]
+        return ["-DENABLE_TENSORFLOW=ON", "-DCMAKE_PREFIX_PATH="+p]
 
     def run(self, extensions):
         print("Building TensorFlow extensions!")