Fix loopy import

firedrakeproject · Oct 23, 2024 · 966aafc · 966aafc
1 parent 09c0be7
commit 966aafc
Showing 1 changed file with 3 additions and 1 deletion.
diff --git a/pyop2/local_kernel.py b/pyop2/local_kernel.py
@@ -4,6 +4,8 @@
 from typing import Union
 
 import loopy as lp
+from loopy.kernel import LoopKernel
+from loopy.translation_unit import TranslationUnit
 from loopy.tools import LoopyKeyBuilder
 import numpy as np
 
@@ -212,7 +214,7 @@ class LoopyLocalKernel(LocalKernel):
         or :class:`loopy.TranslationUnit`.
     """
 
-    @validate_type(("code", (lp.LoopKernel, lp.TranslationUnit), TypeError))
+    @validate_type(("code", (LoopKernel, TranslationUnit), TypeError))
     def __init__(self, code, *args, **kwargs):
         super().__init__(code, *args, **kwargs)