Replace usages of math.exp2

Signed-off-by: Nathan Goldbaum <nathan.goldbaum@gmail.com>
Point72 · Mar 13, 2024 · ebfde8d · ebfde8d
1 parent d276ede
commit ebfde8d
Show file tree

Hide file tree

Showing 2 changed files with 6 additions and 2 deletions.
diff --git a/csp/math.py b/csp/math.py
@@ -1,5 +1,6 @@
 import math
 import numpy as np
+import sys
 import typing
 from functools import lru_cache
 
@@ -363,7 +364,10 @@ def comp(x: ts["T"]):
 log2 = define_unary_op("log2", lambda x: math.log2(x))
 log10 = define_unary_op("log10", lambda x: math.log10(x))
 exp = define_unary_op("exp", lambda x: math.exp(x))
-exp2 = define_unary_op("exp2", lambda x: math.exp2(x))
+if sys.version_info < (3, 11):
+    exp2 = define_unary_op("exp2", lambda x: 2**x)
+else:
+    exp2 = define_unary_op("exp2", lambda x: math.exp2)
 sqrt = define_unary_op("sqrt", lambda x: math.sqrt(x))
 erf = define_unary_op("erf", lambda x: math.erf(x))
 sin = define_unary_op("sin", lambda x: math.sin(x))

diff --git a/csp/tests/test_math.py b/csp/tests/test_math.py
@@ -122,7 +122,7 @@ def test_math_unary_ops(self):
             csp.log2: lambda x: math.log2(x),
             csp.log10: lambda x: math.log10(x),
             csp.exp: lambda x: math.exp(x),
-            csp.exp2: lambda x: math.exp2(x),
+            csp.exp2: lambda x: 2**x,
             csp.sin: lambda x: math.sin(x),
             csp.cos: lambda x: math.cos(x),
             csp.tan: lambda x: math.tan(x),