[Error] Shorten traceback for _BoundedDifferentiableMethod

taichi-dev · Oct 31, 2022 · a96df4d · a96df4d
1 parent b1b7fc8
commit a96df4d
Showing 1 changed file with 6 additions and 3 deletions.
diff --git a/python/taichi/lang/kernel_impl.py b/python/taichi/lang/kernel_impl.py
@@ -1017,9 +1017,12 @@ def __init__(self, kernel_owner, wrapped_kernel_func):
         self.__name__ = None
 
     def __call__(self, *args, **kwargs):
-        if self._is_staticmethod:
-            return self._primal(*args, **kwargs)
-        return self._primal(self._kernel_owner, *args, **kwargs)
+        try:
+            if self._is_staticmethod:
+                return self._primal(*args, **kwargs)
+            return self._primal(self._kernel_owner, *args, **kwargs)
+        except (TaichiCompilationError, TaichiRuntimeError) as e:
+            raise type(e)('\n' + str(e)) from None
 
     def grad(self, *args, **kwargs):
         return self._adjoint(self._kernel_owner, *args, **kwargs)