allow nested set_cufft_callbacks
@@ -410,9 +410,10 @@ def set_cufft_callbacks(
"""
cdef _ThreadLocal tls = _ThreadLocal.get()
- cdef _CallbackManager mgr
+ cdef _CallbackManager mgr, mgr_prev
try:
+ mgr_prev = tls._current_cufft_callback
mgr = _CallbackManager(
cb_load=cb_load,
cb_store=cb_store,
@@ -423,6 +424,6 @@ def set_cufft_callbacks(
_callback_mgr.append(mgr) # keep the manager alive
yield mgr
finally:
- tls._current_cufft_callback = None
+ tls._current_cufft_callback = mgr_prev
# do not remove mgr from _callback_mgr, as one might still wanna use
- # it in the same Python session
+ # plans generated by it in the same Python session
