[bugfix] libentry cache_key supports integer dns arguments with dynam…

…ic value
FlagOpen · Jul 1, 2024 · 8d76d29 · 8d76d29
1 parent aa64c01
commit 8d76d29
Showing 1 changed file with 7 additions and 0 deletions.
diff --git a/src/flag_gems/utils/libentry.py b/src/flag_gems/utils/libentry.py
@@ -38,6 +38,13 @@ def key(self, spec_args, dns_args, const_args):
         for arg in dns_args:
             if hasattr(arg, "data_ptr"):
                 entry_key.append(str(arg.dtype))
+            elif isinstance(arg, int):
+                if -(2**31) <= arg and arg <= 2**31 - 1:
+                    entry_key.append("i32")
+                elif 2**63 <= arg and arg <= 2**64 - 1:
+                    entry_key.append("u64")
+                else:
+                    entry_key.append("i64")
             else:
                 entry_key.append(type(arg))
         # const args passed by position