Skip to content

Commit 562550e

Browse files
committed
revise test function to use fastmath module
1 parent 816e8d4 commit 562550e

File tree

1 file changed

+5
-12
lines changed

1 file changed

+5
-12
lines changed

tests/test_precision.py

Lines changed: 5 additions & 12 deletions
Original file line numberDiff line numberDiff line change
@@ -9,7 +9,7 @@
99
from numba import cuda
1010

1111
import stumpy
12-
from stumpy import cache, config, core
12+
from stumpy import cache, config, core, fastmath
1313

1414
try:
1515
from numba.errors import NumbaPerformanceWarning
@@ -153,14 +153,10 @@ def test_snippets():
153153
): # pragma: no cover
154154
# Revise fastmath flags by removing reassoc (to improve precision),
155155
# recompile njit functions, and re-compute snippets.
156-
core._calculate_squared_distance.targetoptions["fastmath"] = {
157-
"nsz",
158-
"arcp",
159-
"contract",
160-
"afn",
161-
}
156+
fastmath.set_flag(
157+
"core", "_calculate_squared_distance", {"nsz", "arcp", "contract", "afn"}
158+
)
162159
cache._recompile()
163-
164160
(
165161
cmp_snippets,
166162
cmp_indices,
@@ -189,10 +185,7 @@ def test_snippets():
189185

190186
if not numba.config.DISABLE_JIT: # pragma: no cover
191187
# Revert fastmath flag back to their default values
192-
config._reset("STUMPY_FASTMATH_FLAGS")
193-
core._calculate_squared_distance.targetoptions["fastmath"] = (
194-
config.STUMPY_FASTMATH_FLAGS
195-
)
188+
fastmath.set_flag("core", "_calculate_squared_distance")
196189
cache._recompile()
197190

198191

0 commit comments

Comments
 (0)