Skip to content

Commit 95792de

Browse files
authored
[triton_kernels] cacheable sets __name__ and __module__ on JITFunction (triton-lang#7916)
Currently `kernel.__module__` points to the fake module created by `specialize` which means you can't import it programmatically.
1 parent f1872ed commit 95792de

File tree

1 file changed

+3
-0
lines changed

1 file changed

+3
-0
lines changed

python/triton_kernels/triton_kernels/specialize.py

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -19,6 +19,9 @@ def my_kernel(): return (expression dynamically defining a kernel)
1919
g.fn.__name__ = f.__name__
2020
g.fn.__module__ = f.__module__
2121
g.fn.__qualname__ = f.__qualname__
22+
g.__name__ = f.__name__
23+
g.__module__ = f.__module__
24+
g.__qualname__ = f.__qualname__
2225
g._fn_name = f"{f.__module__}.{f.__qualname__}"
2326
return g
2427

0 commit comments

Comments
 (0)