Kernel init graceful failure

hippynn/custom_kernels/env_pytorch.py

@@ -170,7 +170,7 @@ def _featsum_legacy(env, sense, pair_first, pair_second):
     compiler=torch.jit.script,
 )
 
-old_kernels_compile = MessagePassingKernels(
+old_kernels_compile = MessagePassingKernels.safe_init(
     "_legacy_compile",
     _envsum_legacy, _sensesum_legacy, _featsum_legacy,
     compiler=torch.compile,
@@ -197,11 +197,8 @@ def _featsum_legacy(env, sense, pair_first, pair_second):
     compiler=torch.jit.script,
 )
 
-pytorch_kernels_compile = MessagePassingKernels(
+pytorch_kernels_compile = MessagePassingKernels.safe_init(
     "pytorch",
     envsum, sensesum, featsum,
     compiler=torch.compile,
 )
-
-
-

hippynn/custom_kernels/registry.py

@@ -1,3 +1,5 @@
+import warnings
+
 from hippynn.custom_kernels.autograd_wrapper import wrap_envops
 
 
@@ -37,6 +39,20 @@ def __init__(self, impl_name: str, envsum_impl, sensesum_impl, featsum_impl, wra
         else:
             self._registered_implementations[impl_name] = self
 
+    @classmethod
+    def safe_init(self, impl_name, *args, **kwargs):
+        """Init self, but return None if init fails and raise a warning"""
+        try:
+            out = self.__init__(impl_name, *args, **kwargs)
+        except Exception as e:
+            w = RuntimeWarning(*e.args)
+            w.with_traceback(e.__traceback__)
+            warnings.warn(w)
+            out = None
+
+        self._registered_implementations[impl_name] = out
+        return out
+
     @classmethod
     def get_implementation(cls, impl_name):
         """