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;
 }