backport unwrap_if_constexpr

Author: iclementine

src/flag_gems/ops/sort.py

@@ -11,18 +11,26 @@
 logger = logging.getLogger(__name__)
 
 
+def unwrap_if_constexpr(o):
+    return o.value if isinstance(o, tl.constexpr) else o
+
+
 @tl.constexpr
-def get_int_t(num_bits, signed) -> tl.dtype:
+def get_int_t(num_bits: tl.constexpr, signed: tl.constexpr) -> tl.dtype:
+    num_bits = unwrap_if_constexpr(num_bits)
+    signed = unwrap_if_constexpr(signed)
     return tl.core.get_int_dtype(num_bits, signed)
 
 
 @tl.constexpr
-def one_zeros(num_bits) -> int:
+def one_zeros(num_bits: tl.constexpr) -> int:
+    num_bits = unwrap_if_constexpr(num_bits)
     return 1 << (num_bits - 1)
 
 
 @tl.constexpr
-def zero_ones(num_bits) -> int:
+def zero_ones(num_bits: tl.constexpr) -> int:
+    num_bits = unwrap_if_constexpr(num_bits)
     return (1 << (num_bits - 1)) - 1