expand signatures in sjit

pleiszenburg · Dec 29, 2023 · 3baa744 · 3baa744
1 parent 11ff8df
commit 3baa744
Showing 1 changed file with 3 additions and 0 deletions.
diff --git a/src/hapsira/core/jit.py b/src/hapsira/core/jit.py
@@ -160,6 +160,9 @@ def sjit(*args, **kwargs) -> Callable:
     else:
         pass
 
+    if len(args) > 0 and isinstance(args[0], str):
+        args = _parse_signatures(args[0]), *args[1:]
+
     def wrapper(inner_func: Callable) -> Callable:
         """
         Applies JIT