feat: add math.exp2

Author: kaeun97

numba_cuda/numba/cuda/bf16.py

@@ -1,5 +1,6 @@
 # SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 # SPDX-License-Identifier: BSD-2-Clause
+import sys
 
 from numba.cuda._internal.cuda_bf16 import (
     typing_registry,
@@ -191,14 +192,10 @@ def exp_ol(a):
     return _make_unary(a, hexp)
 
 
-try:
-    from math import exp2
-
-    @overload(exp2, target="cuda")
+if sys.version_info >= (3, 11):
+    @overload(math.exp2, target="cuda")
     def exp2_ol(a):
         return _make_unary(a, hexp2)
-except ImportError:
-    pass
 
 ## Public aliases using Numba/Numpy-style type names
 # Floating-point

numba_cuda/numba/cuda/cudamath.py

@@ -1,6 +1,7 @@
 # SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 # SPDX-License-Identifier: BSD-2-Clause
 
+import sys
 import math
 from numba.core import types
 from numba.cuda.typing.templates import ConcreteTemplate, signature, Registry
@@ -57,6 +58,9 @@ class Math_unary_with_fp16(ConcreteTemplate):
         signature(types.float16, types.float16),
     ]
 
+if sys.version_info >= (3, 11):
+    Math_unary_with_fp16 = infer_global(math.exp2)(Math_unary_with_fp16)
+
 
 @infer_global(math.atan2)
 class Math_atan2(ConcreteTemplate):

numba_cuda/numba/cuda/fp16.py

@@ -185,7 +185,7 @@ def log10_ol(a):
     return _make_unary(a, hlog10)
 
 
-@overload(math.exp, target="cuda")
+@overload(, target="cuda")
 def exp_ol(a):
     return _make_unary(a, hexp)
 

numba_cuda/numba/cuda/mathimpl.py

@@ -1,6 +1,7 @@
 # SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 # SPDX-License-Identifier: BSD-2-Clause
 
+import sys
 import math
 import operator
 from llvmlite import ir
@@ -25,6 +26,8 @@
 unarys += [("floor", "floorf", math.floor)]
 unarys += [("fabs", "fabsf", math.fabs)]
 unarys += [("exp", "expf", math.exp)]
+if sys.version_info >= (3, 11):
+    unarys += [("exp2", "exp2f", math.exp2)]
 unarys += [("expm1", "expm1f", math.expm1)]
 unarys += [("erf", "erff", math.erf)]
 unarys += [("erfc", "erfcf", math.erfc)]

numba_cuda/numba/cuda/tests/cudapy/test_bfloat16.py

@@ -1,6 +1,7 @@
 # SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 # SPDX-License-Identifier: BSD-2-Clause
 
+import sys
 import numpy as np
 from ml_dtypes import bfloat16 as mldtypes_bf16
 
@@ -135,12 +136,8 @@ def test_math_bindings(self):
         self.skip_unsupported()
 
         exp_functions = [math.exp]
-        try:
-            from math import exp2
-
-            exp_functions += [exp2]
-        except ImportError:
-            pass
+        if sys.version_info >= (3, 11):
+            exp_functions += [math.exp2]
 
         functions = [
             math.trunc,

numba_cuda/numba/cuda/tests/cudapy/test_math.py

@@ -1,6 +1,7 @@
 # SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 # SPDX-License-Identifier: BSD-2-Clause
 
+import sys
 import numpy as np
 from numba.cuda.testing import (
     skip_unless_cc_53,
@@ -83,6 +84,11 @@ def math_exp(A, B):
     B[i] = math.exp(A[i])
 
 
+def math_exp2(A, B):
+    i = cuda.grid(1)
+    B[i] = math.exp2(A[i])
+
+
 def math_erf(A, B):
     i = cuda.grid(1)
     B[i] = math.erf(A[i])
@@ -400,6 +406,8 @@ def test_math_fp16(self):
         self.unary_template_float16(math_sqrt, np.sqrt)
         self.unary_template_float16(math_ceil, np.ceil)
         self.unary_template_float16(math_floor, np.floor)
+        if sys.version_info >= (3, 11):
+            self.unary_template_float16(math_exp2, np.exp2)
 
     @skip_on_cudasim("numpy does not support trunc for float16")
     @skip_unless_cc_53
@@ -494,6 +502,15 @@ def test_math_exp(self):
         self.unary_template_float64(math_exp, np.exp)
         self.unary_template_int64(math_exp, np.exp)
         self.unary_template_uint64(math_exp, np.exp)
+    # ---------------------------------------------------------------------------
+    # test_math_exp2
+
+    @unittest.skipUnless(sys.version_info >= (3, 11), "Python 3.11+ required")
+    def test_math_exp2(self):
+        self.unary_template_float32(math_exp2, np.exp2)
+        self.unary_template_float64(math_exp2, np.exp2)
+        self.unary_template_int64(math_exp2, np.exp2)
+        self.unary_template_uint64(math_exp2, np.exp2)    
 
     # ---------------------------------------------------------------------------
     # test_math_expm1

numba_cuda/numba/cuda/typing/mathdecl.py

@@ -1,6 +1,7 @@
 # SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 # SPDX-License-Identifier: BSD-2-Clause
 
+import sys
 import math
 from numba.core import types
 from numba.cuda.typing.templates import ConcreteTemplate, signature, Registry
@@ -43,6 +44,8 @@ class Math_unary(ConcreteTemplate):
         signature(types.float64, types.float64),
     ]
 
+if sys.version_info >= (3, 11):
+    Math_unary = infer_global(math.exp2)(Math_unary)
 
 @infer_global(math.atan2)
 class Math_atan2(ConcreteTemplate):