lazy load torch

pykeops/pykeops/__init__.py

@@ -98,10 +98,10 @@ def get_build_folder():
 
 
 if numpy_found:
-    from .numpy.test_install import test_numpy_bindings
+    from ._test_install.numpy import test_numpy_bindings
 
 if torch_found:
-    from .torch.test_install import test_torch_bindings
+    from ._test_install.torch import test_torch_bindings
 
 # next line is to ensure that cache file for formulas is loaded at import
 from .common import keops_io

pykeops/pykeops/torch/test_install.py → pykeops/pykeops/_test_install/torch.py

@@ -1,5 +1,3 @@
-import torch
-
 from pykeops.common.utils import pyKeOps_Message
 
 formula = "SqNorm2(x - y)"
@@ -11,6 +9,9 @@ def test_torch_bindings():
     """
     This function try to compile a simple keops formula using the pytorch binder.
     """
+
+    import torch
+
     x = torch.arange(1, 10, dtype=torch.float32).view(-1, 3)
     y = torch.arange(3, 9, dtype=torch.float32).view(-1, 3)
 

pykeops/setup.py

@@ -55,6 +55,7 @@
         "pykeops.torch.cluster",
         "pykeops.torch.generic",
         "pykeops.torch.lazytensor",
+        "pykeops._test_install",
     ],
     package_data={
         "pykeops": [