Added PyUnstable_Module_SetGIL to PyInit_cpu_feature_guard

jax-ml · Aug 19, 2024 · c591bba · c591bba
1 parent b1b3ea2
commit c591bba
Showing 1 changed file with 8 additions and 1 deletion.
diff --git a/jaxlib/cpu_feature_guard.c b/jaxlib/cpu_feature_guard.c
@@ -172,5 +172,12 @@ static struct PyModuleDef cpu_feature_guard_module = {
 #endif
 
 EXPORT_SYMBOL PyMODINIT_FUNC PyInit_cpu_feature_guard(void) {
- return PyModule_Create(&cpu_feature_guard_module);
+ PyObject *module = PyModule_Create(&cpu_feature_guard_module);
+ if (module == NULL) {
+ return NULL;
+ }
+#ifdef Py_GIL_DISABLED
+ PyUnstable_Module_SetGIL(module, Py_MOD_GIL_NOT_USED);
+#endif
+ return module;
 }