[Backend] Turn to official dlpack for Tensorflow (#1511)

* Turn to official dlpack

* fix

* fix

Jenkinsfile

@@ -102,6 +102,7 @@ pipeline {
         docker {
           label "linux-c52x-node"
           image "dgllib/dgl-ci-lint"  
+          alwaysPull true
         }
       }
       steps {
@@ -121,6 +122,7 @@ pipeline {
             docker {
               label "linux-c52x-node"
               image "dgllib/dgl-ci-cpu:conda" 
+              alwaysPull true
             }
           }
           steps {
@@ -138,6 +140,7 @@ pipeline {
               label "linux-c52x-node"
               image "dgllib/dgl-ci-gpu:conda"
               args "-u root"
+              alwaysPull true
             }
           }
           steps {
@@ -173,6 +176,7 @@ pipeline {
             docker { 
               label "linux-c52x-node"
               image "dgllib/dgl-ci-cpu:conda"
+              alwaysPull true
             }
           }
           steps {
@@ -200,6 +204,7 @@ pipeline {
             docker {
               label "linux-c52x-node"
               image "dgllib/dgl-ci-cpu:conda" 
+              alwaysPull true
             }
           }
           stages {
@@ -221,6 +226,7 @@ pipeline {
               label "linux-gpu-node"
               image "dgllib/dgl-ci-gpu:conda" 
               args "--runtime nvidia"
+              alwaysPull true
             }
           }
           stages {
@@ -241,6 +247,7 @@ pipeline {
             docker {
               label "linux-c52x-node"
               image "dgllib/dgl-ci-cpu:conda" 
+              alwaysPull true
             }
           }
           stages {
@@ -292,6 +299,7 @@ pipeline {
               label "linux-gpu-node"
               image "dgllib/dgl-ci-gpu:conda"
               args "--runtime nvidia"
+              alwaysPull true
             }
           }
           stages {
@@ -318,6 +326,7 @@ pipeline {
             docker {
               label "linux-c52x-node"
               image "dgllib/dgl-ci-cpu:conda" 
+              alwaysPull true
             }
           }
           stages {
@@ -344,6 +353,7 @@ pipeline {
               label "linux-gpu-node" 
               image "dgllib/dgl-ci-gpu:conda"
               args "--runtime nvidia"
+              alwaysPull true
             }
           }
           stages {

docker/install/conda_env/tensorflow_cpu.yml

@@ -3,8 +3,7 @@ dependencies:
   - python=3.6.9
   - pip
   - pip:
-    - tensorflow==2.2.0rc1
-    # - tf-nightly==2.2.0.dev20200327
+    - tensorflow==2.2.0
     - tfdlpack
     - pytest
     - nose

docker/install/conda_env/tensorflow_gpu.yml

@@ -3,8 +3,7 @@ dependencies:
   - python=3.6.9
   - pip
   - pip:
-    - tensorflow==2.2.0rc1
-    # - tf-nightly==2.2.0.dev20200327
+    - tensorflow==2.2.0
     - tfdlpack-gpu
     - pytest
     - nose

python/dgl/backend/tensorflow/tensor.py

@@ -13,7 +13,7 @@ from ... import ndarray as nd
 from ... import kernel as K
 from ...function.base import TargetCode
 
-if os.getenv("USE_OFFICIAL_TFDLPACK", False):
+if not os.getenv("USE_TFDLPACK", False):
     if LooseVersion(tf.__version__) < LooseVersion("2.2.0"):
         raise RuntimeError("DGL requires tensorflow>=2.2.0 for the official DLPack support.")
 
@@ -21,7 +21,7 @@ if os.getenv("USE_OFFICIAL_TFDLPACK", False):
         return tf.experimental.dlpack.to_dlpack(input)
 
     def zerocopy_from_dlpack(dlpack_tensor):
-        # TODO(Jinjing): Tensorflow requires memory to be 64-bit aligned. We check the
+        # TODO(Jinjing): Tensorflow requires memory to be 64-bytes aligned. We check the
         #   alignment and make a copy if needed. The functionality is better in TF's main repo.
         aligned = nd.from_dlpack(dlpack_tensor).to_dlpack(64)
         return tf.experimental.dlpack.from_dlpack(aligned)

tests/scripts/build_dgl.sh

@@ -22,7 +22,7 @@ rm -rf _download
 
 pushd build
 cmake $CMAKE_VARS ..
-make -j4
+make -j8
 popd
 
 pushd python