Repository: dmlc/dgl

Issue:
dgl.knn_graph() error on tensorflow
x = torch.tensor([[0.0, 0.0, 1.0],
                  [1.0, 0.5, 0.5],
                  [0.5, 0.2, 0.2],
                  [0.3, 0.2, 0.4]])
knn_g = dgl.knn_graph(x, 2)

There is no problem on PyTorch  

But on Tensorflow:
x = tf.constant([[0.0, 0.0, 1.0],
                  [1.0, 0.5, 0.5],
                  [0.5, 0.2, 0.2],
                  [0.3, 0.2, 0.4]])
knn_g = dgl.knn_graph(x, 2)

Traceback (most recent call last):
  File "<input>", line 5, in <module>
  File "D:\E\Anaconda3\envs\tfgpu\lib\site-packages\dgl\transform.py", line 187, in knn_graph
    return _knn_graph_blas(x, k, dist=dist)
  File "D:\E\Anaconda3\envs\tfgpu\lib\site-packages\dgl\transform.py", line 240, in _knn_graph_blas
    dist = pairwise_squared_distance(x)
  File "D:\E\Anaconda3\envs\tfgpu\lib\site-packages\dgl\transform.py", line 67, in pairwise_squared_distance
    return x2s + F.swapaxes(x2s, -1, -2) - 2 * x @ F.swapaxes(x, -1, -2)
  File "D:\E\Anaconda3\envs\tfgpu\lib\site-packages\dgl\backend\tensorflow\tensor.py", line 323, in swapaxes
    return tf.transpose(input, perm=[axis1, axis2])
  File "D:\E\Anaconda3\envs\tfgpu\lib\site-packages\tensorflow\python\util\dispatch.py", line 201, in wrapper
    return target(*args, **kwargs)
  File "D:\E\Anaconda3\envs\tfgpu\lib\site-packages\tensorflow\python\ops\array_ops.py", line 2135, in transpose_v2
    return transpose(a=a, perm=perm, name=name, conjugate=conjugate)
  File "D:\E\Anaconda3\envs\tfgpu\lib\site-packages\tensorflow\python\util\dispatch.py", line 201, in wrapper
    return target(*args, **kwargs)
  File "D:\E\Anaconda3\envs\tfgpu\lib\site-packages\tensorflow\python\ops\array_ops.py", line 2216, in transpose
    return transpose_fn(a, perm, name=name)
  File "D:\E\Anaconda3\envs\tfgpu\lib\site-packages\tensorflow\python\ops\gen_array_ops.py", line 11589, in transpose
    x, perm, name=name, ctx=_ctx)
  File "D:\E\Anaconda3\envs\tfgpu\lib\site-packages\tensorflow\python\ops\gen_array_ops.py", line 11614, in transpose_eager_fallback
    attrs=_attrs, ctx=ctx, name=name)
  File "D:\E\Anaconda3\envs\tfgpu\lib\site-packages\tensorflow\python\eager\execute.py", line 60, in quick_execute
    inputs, attrs, num_outputs)
tensorflow.python.framework.errors_impl.InvalidArgumentError: transpose expects a vector of size 3. But input(1) is a vector of size 2 [Op:Transpose]


I do not understand why this problem occurs. I would appreciate receiving your reply


Discussion:

Resolving patch:
diff --git a/python/dgl/backend/tensorflow/tensor.py b/python/dgl/backend/tensorflow/tensor.py
--- a/python/dgl/backend/tensorflow/tensor.py
+++ b/python/dgl/backend/tensorflow/tensor.py
@@ -320,7 +320,10 @@ def reshape(input, shape):
 
 
 def swapaxes(input, axis1, axis2):
-    return tf.transpose(input, perm=[axis1, axis2])
+    ndim = input.ndim
+    t = list(range(ndim))
+    t[axis1], t[axis2] = axis2 % ndim, axis1 % ndim
+    return tf.transpose(input, perm=t)
 
 
 def zeros(shape, dtype, ctx):
diff --git a/python/dgl/transform.py b/python/dgl/transform.py
--- a/python/dgl/transform.py
+++ b/python/dgl/transform.py
@@ -254,7 +254,7 @@ def _knn_graph_blas(x, k, dist='euclidean'):
 
     ctx = F.context(x)
     dist = pairwise_squared_distance(x)
-    k_indices = F.argtopk(dist, k, 2, descending=False)
+    k_indices = F.astype(F.argtopk(dist, k, 2, descending=False), F.int64)
     # index offset for each sample
     offset = F.arange(0, n_samples, ctx=ctx) * n_points
     offset = F.unsqueeze(offset, 1)


Test patch: