refactor into a python package

attempt to load dependent cpu/cuda submodules before resorting to JIT
add pyproject.toml to facilitate package install-time dependency on pytorch and ninja
add setup.py

pyproject.toml

+[build-system]
+requires = ['setuptools>=38.2.5', 'wheel', 'torch>=1.5', 'ninja']
+build-backend = "setuptools.build_meta"
\ No newline at end of file

requirements.txt

 torch>=1.5
-ninja
\ No newline at end of file

setup.py

-from setuptools import setup
+#!/usr/bin/env python
+import os
+from setuptools import setup, find_packages
 from torch.utils.cpp_extension import BuildExtension, CppExtension, CUDAExtension
 
+
+with open('requirements.txt') as f:
+    requirements = f.read().splitlines()
+
+
 setup(
     name='torch_discounted_cumsum',
+    version=0.1,
+    description='Fast differentiable discounted cumulative sum',
+    install_requires=requirements,
+    python_requires='>=3.6',
+    packages=find_packages(),
+    author='Anton Obukhov',
+    license='BSD',
+    url='https://www.github.com/toshas/torch-discounted-cumsum',
     ext_modules=[
-        CppExtension('torch_discounted_cumsum_cpu', [
-            'discounted_cumsum_cpu.cpp'
-        ]),
-        CUDAExtension('torch_discounted_cumsum_cuda', [
-            'discounted_cumsum_cuda.cpp',
-            'discounted_cumsum_cuda_kernel.cu',
-        ])
+        CppExtension(
+            'torch_discounted_cumsum_cpu',
+            [
+                os.path.join('torch_discounted_cumsum', 'discounted_cumsum_cpu.cpp'),
+            ],
+        ),
+        CUDAExtension(
+            'torch_discounted_cumsum_cuda',
+            [
+                os.path.join('torch_discounted_cumsum', 'discounted_cumsum_cuda.cpp'),
+                os.path.join('torch_discounted_cumsum', 'discounted_cumsum_cuda_kernel.cu'),
+            ],
+        )
     ],
     cmdclass={
         'build_ext': BuildExtension
-    })
+    },
+    keywords=[
+        'pytorch', 'discounted', 'cumsum', 'cumulative', 'sum', 'scan', 'differentiable',
+        'reinforcement', 'learning', 'rewards', 'time', 'series'
+    ],
+)

test.py

 import time
 import torch
 
-from discounted_cumsum import discounted_cumsum_left, discounted_cumsum_right
+from torch_discounted_cumsum import discounted_cumsum_left, discounted_cumsum_right
 
 
 def discounted_cumsum_left_gold(input, gamma):

torch_discounted_cumsum/__init__.py

+from .discounted_cumsum import discounted_cumsum_left, discounted_cumsum_right

discounted_cumsum.py → torch_discounted_cumsum/discounted_cumsum.py

+import os
+
 import torch
 from torch.utils.cpp_extension import load
 
+VERBOSE = False
+
+
+def _resolve(name):
+    return os.path.join(os.path.dirname(os.path.realpath(__file__)), name)
 
-torch_discounted_cumsum_cpu = load(
-    name='torch_discounted_cumsum_cpu',
-    sources=['discounted_cumsum_cpu.cpp'],
-    # verbose=True,
-)
 
-torch_discounted_cumsum_cuda = None
-if torch.cuda.is_available():
-    torch_discounted_cumsum_cuda = load(
-        name='torch_discounted_cumsum_cuda',
-        sources=['discounted_cumsum_cuda.cpp', 'discounted_cumsum_cuda_kernel.cu'],
-        verbose=True,
+try:
+    import torch_discounted_cumsum_cpu
+except ImportError:
+    if VERBOSE:
+        print('Falling back to JIT compiling torch_discounted_cumsum_cpu')
+    torch_discounted_cumsum_cpu = load(
+        name='torch_discounted_cumsum_cpu',
+        sources=[
+            _resolve('discounted_cumsum_cpu.cpp'),
+        ],
+        verbose=VERBOSE,
     )
 
 
+try:
+    import torch_discounted_cumsum_cuda
+except ImportError:
+    if VERBOSE:
+        print('Falling back to JIT compiling torch_discounted_cumsum_cuda')
+    torch_discounted_cumsum_cuda = None
+    if torch.cuda.is_available():
+        torch_discounted_cumsum_cuda = load(
+            name='torch_discounted_cumsum_cuda',
+            sources=[
+                _resolve('discounted_cumsum_cuda.cpp'),
+                _resolve('discounted_cumsum_cuda_kernel.cu'),
+            ],
+            verbose=VERBOSE,
+        )
+
+
 def _discounted_cumsum_left_dispatcher(input, gamma):
     if not torch.is_tensor(input):
         raise ValueError('Input must be a torch.Tensor')