better loading

Author: rusty1s

.travis.yml

@@ -77,11 +77,10 @@ install:
 
 script:
   - if [ "${PYTHON_VERSION}" != "3.5" ]; then pip install flake8 && flake8 .; fi
-  - if [ "${PYTHON_VERSION}" = "3.5" ]; then pip install zzip; fi
   - python setup.py test
 after_success:
-  - python setup.py bdist_wheel --dist-dir=dist/torch-${TORCH_VERSION}/${IDX}
-  - python script/rename_wheel.py
+  - python setup.py bdist_wheel --dist-dir=dist/torch-${TORCH_VERSION}
+  - python script/rename_wheel.py ${IDX}
   - pip install codecov && codecov
 deploy:
   provider: s3
@@ -90,8 +89,8 @@ deploy:
   access_key_id: AKIAJB7S6NJ5OM5MAAGA
   secret_access_key: ${S3_SECRET_ACCESS_KEY}
   bucket: pytorch-scatter
-  local_dir: dist/torch-${TORCH_VERSION}/${IDX}
-  upload_dir: whl/torch-${TORCH_VERSION}/${IDX}
+  local_dir: dist/torch-${TORCH_VERSION}
+  upload_dir: whl/torch-${TORCH_VERSION}
   acl: public_read
   on:
     repo: rusty1s/pytorch_scatter

script/rename_wheel.py

@@ -1,19 +1,24 @@
+import sys
 import os
 import os.path as osp
 import glob
+import shutil
+
+idx = sys.argv[1]
+assert idx in ['cpu', 'cu92', 'cu100', 'cu101']
 
 dist_dir = osp.join(osp.dirname(osp.abspath(__file__)), '..', 'dist')
 wheels = glob.glob(osp.join('dist', '**', '*.whl'), recursive=True)
 
 for wheel in wheels:
-    idx = wheel.split(osp.sep)[-2]
-    if idx not in ['cpu', 'cu92', 'cu100', 'cu101']:
-        continue
-    name = wheel.split(osp.sep)[-1]
-    if idx in name:
+    if idx in wheel:
         continue
 
-    names = name.split('-')
-    name = '-'.join(names[:-4] + [names[-4] + '%2B' + idx] + names[-2:])
-    new_wheel = osp.join(*wheel.split(osp.sep)[:-1], name)
-    os.rename(wheel, new_wheel)
+    paths = wheel.split(osp.sep)
+    names = paths[-1].split('-')
+
+    name = '-'.join(names[:-4] + ['latest+' + idx] + names[-3:])
+    shutil.copyfile(wheel, osp.join(*paths[:-1], name))
+
+    name = '-'.join(names[:-4] + [names[-4] + '+' + idx] + names[-3:])
+    os.rename(wheel, osp.join(*paths[:-1], name))

torch_scatter/__init__.py

@@ -1,6 +1,6 @@
 # flake8: noqa
 
-import importlib
+import glob
 import os.path as osp
 
 import torch
@@ -9,8 +9,8 @@
 expected_torch_version = (1, 4)
 
 try:
-    torch.ops.load_library(importlib.machinery.PathFinder().find_spec(
-        '_version', [osp.dirname(__file__)]).origin)
+    torch.ops.load_library(
+        glob.glob(osp.join(osp.dirname(__file__), '_version.*'))[0])
 except OSError as e:
     if 'undefined symbol' in str(e):
         major, minor = [int(x) for x in torch.__version__.split('.')[:2]]

torch_scatter/scatter.py

@@ -1,13 +1,13 @@
-import importlib
+import glob
 import os.path as osp
 from typing import Optional, Tuple
 
 import torch
 
 from .utils import broadcast
 
-torch.ops.load_library(importlib.machinery.PathFinder().find_spec(
-    '_scatter', [osp.dirname(__file__)]).origin)
+torch.ops.load_library(
+    glob.glob(osp.join(osp.dirname(__file__), '_scatter.*'))[0])
 
 
 @torch.jit.script

torch_scatter/segment_coo.py

@@ -1,11 +1,11 @@
-import importlib
+import glob
 import os.path as osp
 from typing import Optional, Tuple
 
 import torch
 
-torch.ops.load_library(importlib.machinery.PathFinder().find_spec(
-    '_segment_coo', [osp.dirname(__file__)]).origin)
+torch.ops.load_library(
+    glob.glob(osp.join(osp.dirname(__file__), '_segment_coo.*'))[0])
 
 
 @torch.jit.script

torch_scatter/segment_csr.py

@@ -1,11 +1,11 @@
-import importlib
+import glob
 import os.path as osp
 from typing import Optional, Tuple
 
 import torch
 
-torch.ops.load_library(importlib.machinery.PathFinder().find_spec(
-    '_segment_csr', [osp.dirname(__file__)]).origin)
+torch.ops.load_library(
+    glob.glob(osp.join(osp.dirname(__file__), '_segment_csr.*'))[0])
 
 
 @torch.jit.script