reorganize import for LazyTensor

getkeops · Jul 11, 2019 · f93e39b · f93e39b
1 parent d5fc83a
commit f93e39b
Show file tree

Hide file tree

Showing 30 changed files with 58 additions and 62 deletions.
diff --git a/doc/index.rst b/doc/index.rst
@@ -23,7 +23,7 @@ Using the **PyTorch backend**, a typical sample of code looks like:
     y = torch.randn(2000000, 3).cuda()
 
     # Turn our Tensors into KeOps symbolic variables:
-    from pykeops import LazyTensor
+    from pykeops.torch import LazyTensor
     x_i = LazyTensor( x[:,None,:] )  # x_i.shape = (1e6, 1, 3)
     y_j = LazyTensor( y[None,:,:] )  # y_j.shape = ( 1, 2e6,3)
 

diff --git a/doc/python/LazyTensor.rst b/doc/python/LazyTensor.rst
@@ -20,7 +20,7 @@ this new tensor type may be used with **very little overhead**:
     y = torch.randn(2000000, 3).cuda()
 
     # Turn our Tensors into KeOps symbolic variables:
-    from pykeops import LazyTensor
+    from pykeops.torch import LazyTensor
     x_i = LazyTensor( x[:,None,:] )  # x_i.shape = (1e6, 1, 3)
     y_j = LazyTensor( y[None,:,:] )  # y_j.shape = ( 1, 2e6,3)
 

diff --git a/pykeops/__init__.py b/pykeops/__init__.py
@@ -24,6 +24,3 @@
 build_type = str(os.environ['PYKEOPS_BUILD_TYPE']) if ('PYKEOPS_BUILD_TYPE' in os.environ) else 'Release'
 
 sys.path.append(bin_folder)
-
-
-from .common.lazy_tensor import LazyTensor, Vi, Vj, Pm
diff --git a/pykeops/benchmarks/plot_benchmarks_convolutions_3D.py b/pykeops/benchmarks/plot_benchmarks_convolutions_3D.py
@@ -118,7 +118,7 @@ def gaussianconv_pytorch(x, y, b):
 #############################################
 # Finally, perform the same operation with our high-level :mod:`LazyTensor <pykeops.common.lazy_tensor.LazyTensor>` wrapper:
 
-from pykeops import LazyTensor
+from pykeops.torch import LazyTensor
 
 def gaussianconv_lazytensor(x, y, b):
     nbatchdims = len(x.shape) - 2

diff --git a/pykeops/common/lazy_tensor.py b/pykeops/common/lazy_tensor.py
@@ -1,19 +1,14 @@
 import numpy as np
-from pykeops.numpy import Genred as Genred_numpy
-from pykeops.numpy import KernelSolve as KernelSolve_numpy
-from pykeops.numpy.utils import numpytools
 import re
 import copy
 
-try:
+from pykeops import torch_found as usetorch
+
+if not usetorch:
+    from pykeops.numpy.utils import numpytools
+else:
     import torch
-    from pykeops.torch import Genred as Genred_torch
-    from pykeops.torch import KernelSolve as KernelSolve_torch
     from pykeops.torch.utils import torchtools
-    usetorch = True
-except ImportError:
-    usetorch = False
-    pass
 
 
 
@@ -170,14 +165,14 @@ def __init__(self, x=None, axis=None):
 
             if typex == np.ndarray:
                 self.tools = numpytools
-                self.Genred = Genred_numpy
-                self.KernelSolve = KernelSolve_numpy
+                self.Genred = numpytools.Genred
+                self.KernelSolve = numpytools.KernelSolve
                 self.dtype = self.tools.dtypename( self.tools.dtype(x) )
 
             elif usetorch and typex == torch.Tensor:
                 self.tools = torchtools
-                self.Genred = Genred_torch
-                self.KernelSolve = KernelSolve_torch
+                self.Genred = torchtools.Genred
+                self.KernelSolve = torchtools.KernelSolve
                 self.dtype = self.tools.dtypename( self.tools.dtype(x) )
             else:
                 raise ValueError("LazyTensors should be built from NumPy arrays, PyTorch tensors, " \
@@ -602,13 +597,13 @@ def __call__(self, *args, **kwargs):
         if self.dtype is None:  # This can only happen if we haven't encountered 2D or 3D arrays just yet...
             if isinstance(args[0], np.ndarray):  # We use the "NumPy" or "PyTorch" backend depending on the first argument
                 self.tools = numpytools
-                self.Genred = Genred_numpy
-                self.KernelSolve = KernelSolve_numpy
+                self.Genred = numpytools.Genred
+                self.KernelSolve = numpytools.KernelSolve
 
             elif usetorch and isinstance(args[0], torch.Tensor):
                 self.tools = torchtools
-                self.Genred = Genred_torch
-                self.KernelSolve = KernelSolve_torch
+                self.Genred = torchtools.Genred
+                self.KernelSolve = torchtools.KernelSolve
 
             self.dtype = self.tools.dtypename( self.tools.dtype(args[0]) )
             self.fixvariables()

diff --git a/pykeops/examples/numpy/gpu_select_numpy_helper.py b/pykeops/examples/numpy/gpu_select_numpy_helper.py
@@ -19,8 +19,8 @@
 import torch
 import matplotlib.pyplot as plt
 
-from pykeops import LazyTensor as keops
-from pykeops import Vi, Vj, Pm 
+from pykeops.torch import LazyTensor as keops
+from pykeops.torch import Vi, Vj, Pm
 from pykeops.numpy.utils import IsGpuAvailable
 
 ###############################################################

diff --git a/pykeops/examples/numpy/plot_grid_cluster_numpy.py b/pykeops/examples/numpy/plot_grid_cluster_numpy.py
@@ -18,7 +18,7 @@
 import numpy as np
 from matplotlib import pyplot as plt
 
-from pykeops import LazyTensor
+from pykeops.numpy import LazyTensor
 from pykeops.numpy.utils import IsGpuAvailable
 
 use_cuda = IsGpuAvailable()

diff --git a/pykeops/examples/numpy/plot_test_invkernel_numpy_helper.py b/pykeops/examples/numpy/plot_test_invkernel_numpy_helper.py
@@ -19,7 +19,7 @@
 import time 
 import matplotlib.pyplot as plt
 
-from pykeops import Vi, Vj, Pm
+from pykeops.numpy import Vi, Vj, Pm
 from pykeops.numpy import KernelSolve
 from pykeops.numpy.utils import IsGpuAvailable
 

diff --git a/pykeops/examples/numpy/plot_test_softmax_numpy_helper.py b/pykeops/examples/numpy/plot_test_softmax_numpy_helper.py
@@ -33,12 +33,13 @@
 
 import time
 import numpy as np
-from pykeops.numpy import Genred
-from pykeops import LazyTensor as kf
-from pykeops import Vi, Vj, Pm
-from pykeops.numpy.utils import WarmUpGpu
 import matplotlib.pyplot as plt
 
+from pykeops.numpy import LazyTensor as kf
+from pykeops.numpy import Vi, Vj, Pm
+from pykeops.numpy.utils import WarmUpGpu
+
+
 ###############################################################################
 # Define our dataset:
 #

diff --git a/pykeops/examples/pytorch/plot_grid_cluster_pytorch.py b/pykeops/examples/pytorch/plot_grid_cluster_pytorch.py
@@ -19,7 +19,7 @@
 import torch
 from matplotlib import pyplot as plt
 
-from pykeops import LazyTensor
+from pykeops.torch import LazyTensor
 
 nump = lambda t : t.cpu().numpy()
 use_cuda = torch.cuda.is_available()

diff --git a/pykeops/examples/pytorch/plot_test_invkernel_torch_helper.py b/pykeops/examples/pytorch/plot_test_invkernel_torch_helper.py
@@ -19,8 +19,8 @@
 
 import torch
 
-from pykeops import Vi, Vj, Pm
-from pykeops import LazyTensor as keops
+from pykeops.torch import Vi, Vj, Pm
+from pykeops.torch import LazyTensor as keops
 
 
 ###############################################################################

diff --git a/pykeops/numpy/__init__.py b/pykeops/numpy/__init__.py
@@ -19,7 +19,8 @@
 from .operations import KernelSolve
 from .convolutions.radial_kernel import RadialKernelConv, RadialKernelGrad1conv
 from .generic.generic_ops import generic_sum, generic_logsumexp, generic_argmin, generic_argkmin
+from pykeops.common.lazy_tensor import LazyTensor, Vi, Vj, Pm
 
-__all__ = sorted(["Genred", "generic_sum", "generic_logsumexp", "generic_argmin", "generic_argkmin", "KernelSolve"])
+__all__ = sorted(["Genred", "generic_sum", "generic_logsumexp", "generic_argmin", "generic_argkmin", "KernelSolve", "LazyTensor", "Vi", "Vj", "Pm"])
 
 
diff --git a/pykeops/numpy/utils.py b/pykeops/numpy/utils.py
@@ -1,6 +1,6 @@
 import numpy as np
 
-from pykeops.numpy import Genred, default_dtype
+from pykeops.numpy import Genred, default_dtype, KernelSolve
 from pykeops.numpy.cluster import swap_axes as np_swap_axes
 
 
@@ -10,7 +10,7 @@ class numpytools:
     exp = np.exp
     log = np.log
     Genred = Genred
-
+    KernelSolve = KernelSolve
     swap_axes = np_swap_axes
 
     @staticmethod

diff --git a/pykeops/test/unit_tests_numpy.py b/pykeops/test/unit_tests_numpy.py
@@ -266,7 +266,7 @@ def test_argkmin(self):
     ############################################################
     def test_LazyTensor_sum(self):
     ############################################################
-        from pykeops import LazyTensor
+        from pykeops.numpy import LazyTensor
 
         full_results = []
         for use_keops in [True, False]:

diff --git a/pykeops/test/unit_tests_pytorch.py b/pykeops/test/unit_tests_pytorch.py
@@ -401,7 +401,7 @@ def test_pickle(self):
     def test_LazyTensor_sum(self):
     ############################################################
         import torch
-        from pykeops import LazyTensor
+        from pykeops.torch import LazyTensor
 
         full_results = []
         for use_keops in [True, False]:
@@ -441,7 +441,7 @@ def test_LazyTensor_sum(self):
     def test_LazyTensor_logsumexp(self):
     ############################################################
         import torch
-        from pykeops import LazyTensor
+        from pykeops.torch import LazyTensor
 
         full_results = []
         for use_keops in [True, False]:
@@ -480,7 +480,7 @@ def test_LazyTensor_logsumexp(self):
     def test_LazyTensor_min(self):
     ############################################################
         import torch
-        from pykeops import LazyTensor
+        from pykeops.torch import LazyTensor
 
         full_results = []
         for use_keops in [True, False]:

diff --git a/pykeops/torch/__init__.py b/pykeops/torch/__init__.py
@@ -27,5 +27,6 @@
 from .kernel_product.kernels import Kernel, kernel_product, kernel_formulas
 from .generic.generic_ops import generic_sum, generic_logsumexp, generic_argmin, generic_argkmin
 from .kernel_product.formula import Formula
+from pykeops.common.lazy_tensor import LazyTensor, Vi, Vj, Pm
 
-__all__ = sorted(["Genred", "generic_sum", "generic_logsumexp", "generic_argmin", "generic_argkmin", "Kernel", "kernel_product", "KernelSolve", "kernel_formulas", "Formula"])
+__all__ = sorted(["Genred", "generic_sum", "generic_logsumexp", "generic_argmin", "generic_argkmin", "Kernel", "kernel_product", "KernelSolve", "kernel_formulas", "Formula", "LazyTensor", "Vi", "Vj", "Pm"])
diff --git a/pykeops/torch/utils.py b/pykeops/torch/utils.py
@@ -1,6 +1,6 @@
 import torch
 
-from pykeops.torch import Genred, default_dtype
+from pykeops.torch import Genred, KernelSolve, default_dtype
 from pykeops.torch.cluster import swap_axes as torch_swap_axes
 
 
@@ -19,7 +19,8 @@ class torchtools:
 
     swap_axes = torch_swap_axes
 
-    # Genred = Genred
+    Genred = Genred
+    KernelSolve = KernelSolve
     # GenredLowlevel = GenredLowlevel
 
     @staticmethod

diff --git a/pykeops/tutorials/a_LazyTensors/plot_lazytensors_a.py b/pykeops/tutorials/a_LazyTensors/plot_lazytensors_a.py
@@ -107,7 +107,7 @@
 # 
 # 
 
-from pykeops import LazyTensor
+from pykeops.torch import LazyTensor
 
 x_i = LazyTensor( x[:,None,:] )  # (M, 1, 2) KeOps LazyTensor, wrapped around the numpy array x
 y_j = LazyTensor( y[None,:,:] )  # (1, N, 2) KeOps LazyTensor, wrapped around the numpy array y

diff --git a/pykeops/tutorials/a_LazyTensors/plot_lazytensors_b.py b/pykeops/tutorials/a_LazyTensors/plot_lazytensors_b.py
@@ -27,7 +27,7 @@
 # formulas:
 
 import torch
-from pykeops import LazyTensor
+from pykeops.torch import LazyTensor
 use_cuda = torch.cuda.is_available()
 tensor = torch.cuda.FloatTensor if use_cuda else torch.FloatTensor
 M, N = (100000, 200000) if use_cuda else (1000, 2000)

diff --git a/pykeops/tutorials/a_LazyTensors/plot_lazytensors_c.py b/pykeops/tutorials/a_LazyTensors/plot_lazytensors_c.py
@@ -9,7 +9,7 @@
 
 import time
 import torch
-from pykeops import LazyTensor
+from pykeops.torch import LazyTensor
 use_cuda = torch.cuda.is_available()
 tensor = torch.cuda.FloatTensor if use_cuda else torch.FloatTensor
 
@@ -36,7 +36,7 @@
 
 ############################################################################
 # Then we use Vi and Vj to convert to KeOps LazyTensor objects
-from pykeops import Vi, Vj, Pm
+from pykeops.torch import Vi, Vj, Pm
 x_i = Vi(x)  # (M, 1, D) LazyTensor, equivalent to LazyTensor( x[:,None,:] ) 
 y_j = Vj(y)  # (1, N, D) LazyTensor, equivalent to LazyTensor( y[None,:,:] ) 
 
@@ -138,7 +138,7 @@
 y = torch.randn(N, D).type(tensor)
 b = torch.randn(N, Dv).type(tensor)
 sigmas = tensor([0.5, 1.0, 2.0, 4.0])
-from pykeops import Vi, Vj, Pm
+from pykeops.torch import Vi, Vj, Pm
 x_i = Vi(x)  # (M, 1, D) LazyTensor, equivalent to LazyTensor( x[:,None,:] ) 
 y_j = Vj(y)  # (1, N, D) LazyTensor, equivalent to LazyTensor( y[None,:,:] ) 
 b_j = Vj(b)  # (1, N, D) LazyTensor, equivalent to LazyTensor( b[None,:,:] ) 

diff --git a/pykeops/tutorials/backends/plot_gpytorch.py b/pykeops/tutorials/backends/plot_gpytorch.py
@@ -66,7 +66,7 @@
 #   we hope to provide a simpler interface in future releases. 
 
 
-from pykeops import LazyTensor
+from pykeops.torch import LazyTensor
 
 class KeOpsRBFLazyTensor(gpytorch.lazy.LazyTensor):
     def __init__(self, x_i, y_j):

diff --git a/pykeops/tutorials/backends/plot_scipy.py b/pykeops/tutorials/backends/plot_scipy.py
@@ -34,7 +34,8 @@
 import matplotlib.pyplot as plt
 import numpy as np
 
-from pykeops import LazyTensor
+#from pykeops.common.lazy_tensor import LazyTensor
+from pykeops.numpy import LazyTensor
 from pykeops.numpy.utils import IsGpuAvailable
 
 use_cuda = IsGpuAvailable()

diff --git a/pykeops/tutorials/interpolation/plot_RBF_interpolation_numpy.py b/pykeops/tutorials/interpolation/plot_RBF_interpolation_numpy.py
@@ -32,7 +32,7 @@
 import numpy as np
 from matplotlib import pyplot as plt
 
-from pykeops import LazyTensor
+from pykeops.numpy import LazyTensor
 from pykeops.numpy.utils import IsGpuAvailable
 
 #######################################################################

diff --git a/pykeops/tutorials/interpolation/plot_RBF_interpolation_torch.py b/pykeops/tutorials/interpolation/plot_RBF_interpolation_torch.py
@@ -32,7 +32,7 @@
 import torch
 from matplotlib import pyplot as plt
 
-from pykeops import LazyTensor
+from pykeops.torch import LazyTensor
 
 ###############################################################################################
 # Generate some data:

diff --git a/pykeops/tutorials/kmeans/plot_kmeans_numpy.py b/pykeops/tutorials/kmeans/plot_kmeans_numpy.py
@@ -28,7 +28,7 @@
 import numpy as np
 from matplotlib import pyplot as plt
 
-from pykeops import LazyTensor
+from pykeops.numpy import LazyTensor
 from pykeops.numpy.utils import IsGpuAvailable
 
 dtype = 'float32'  # May be 'float32' or 'float64'

diff --git a/pykeops/tutorials/kmeans/plot_kmeans_torch.py b/pykeops/tutorials/kmeans/plot_kmeans_torch.py
@@ -27,7 +27,7 @@
 import torch
 from matplotlib import pyplot as plt
 
-from pykeops import LazyTensor
+from pykeops.torch import LazyTensor
 
 use_cuda = torch.cuda.is_available()
 dtype = 'float32' if use_cuda else 'float64'

diff --git a/pykeops/tutorials/knn/plot_knn_mnist.py b/pykeops/tutorials/knn/plot_knn_mnist.py
@@ -22,10 +22,9 @@
 import time
 from matplotlib import pyplot as plt
 
-import numpy as np
 import torch
 
-from pykeops import LazyTensor
+from pykeops.torch import LazyTensor
 
 use_cuda = torch.cuda.is_available()
 tensor = torch.cuda.FloatTensor if use_cuda else torch.FloatTensor

diff --git a/pykeops/tutorials/knn/plot_knn_numpy.py b/pykeops/tutorials/knn/plot_knn_numpy.py
@@ -20,7 +20,7 @@
 
 import time
 import numpy as np
-from pykeops import LazyTensor
+from pykeops.numpy import LazyTensor
 from pykeops.numpy.utils import IsGpuAvailable
 from matplotlib import pyplot as plt
 

diff --git a/pykeops/tutorials/knn/plot_knn_torch.py b/pykeops/tutorials/knn/plot_knn_torch.py
@@ -24,7 +24,7 @@
 import numpy as np
 import torch
 
-from pykeops import LazyTensor
+from pykeops.torch import LazyTensor
 
 use_cuda = torch.cuda.is_available()
 dtype = torch.cuda.FloatTensor if use_cuda else torch.FloatTensor

diff --git a/readme.md b/readme.md
@@ -21,7 +21,7 @@ x = torch.randn(1000000, 3, requires_grad=True).cuda()
 y = torch.randn(2000000, 3).cuda()
 
 # Turn our Tensors into KeOps symbolic variables:
-from pykeops import LazyTensor
+from pykeops.torch import LazyTensor
 x_i = LazyTensor( x[:,None,:] )  # x_i.shape = (1e6, 1, 3)
 y_j = LazyTensor( y[None,:,:] )  # y_j.shape = ( 1, 2e6,3)
 
@@ -79,6 +79,6 @@ As of today, KeOps provides core routines for:
 
 Feel free to contact us for any bug report or feature request:
 
-* [Benjamin Charlier](http://imag.umontpellier.fr/~charlier/)
+* [Benjamin Charlier](https://imag.umontpellier.fr/~charlier/)
 * [Jean Feydy](https://www.math.ens.fr/~feydy/)
 * [Joan Alexis Glaunès](https://www.mi.parisdescartes.fr/~glaunes/)