update kernel name and debug info (FlagOpen#444)

Co-authored-by: kiddyjinjin <kiddyjinjin@gmailcom>
tianxiao-baai · Feb 13, 2025 · de58cc7 · de58cc7
1 parent a09d014
commit de58cc7
Show file tree

Hide file tree

Showing 4 changed files with 4 additions and 4 deletions.
diff --git a/src/flag_gems/ops/gather.py b/src/flag_gems/ops/gather.py
@@ -199,7 +199,7 @@ def __call__(self, *args, **kwargs):
             code = generate_code(
                 args,
                 "_gather_wrapper",
-                "_gather_jit_function",
+                "_gather_flaggems_jit_function",
                 code,
             )
 

diff --git a/src/flag_gems/ops/repeat.py b/src/flag_gems/ops/repeat.py
@@ -404,7 +404,7 @@ def __call__(self, x, sizes):
                 ndim,
                 "_wrapper",
                 "_wrapper_out",
-                "_jit_function",
+                "_repeat_flaggems_jit_function",
                 code,
             )
 

diff --git a/src/flag_gems/ops/tile.py b/src/flag_gems/ops/tile.py
@@ -404,7 +404,7 @@ def __call__(self, x, dims):
                 ndim,
                 "_wrapper",
                 "_wrapper_out",
-                "_jit_function",
+                "_tile_flaggems_jit_function",
                 code,
             )
 

diff --git a/src/flag_gems/ops/zeros_like.py b/src/flag_gems/ops/zeros_like.py
@@ -10,7 +10,7 @@
 def zeros_like(
     x, *, dtype=None, layout=None, device=None, pin_memory=None, memory_format=None
 ):
-    logging.debug("GEMS FULL_LIKE")
+    logging.debug("GEMS ZEROS_LIKE")
     if device is None:
         device = x.device
     if dtype is None: