From c591bba95f07cf2a411d55de617428b7db731d78 Mon Sep 17 00:00:00 2001
From: vfdev-5 <vfdev.5@gmail.com>
Date: Tue, 20 Aug 2024 00:10:17 +0200
Subject: [PATCH] Added PyUnstable_Module_SetGIL to PyInit_cpu_feature_guard

---
 jaxlib/cpu_feature_guard.c | 9 ++++++++-
 1 file changed, 8 insertions(+), 1 deletion(-)

diff --git a/jaxlib/cpu_feature_guard.c b/jaxlib/cpu_feature_guard.c
index 7c8ff2951a79..d18478eb57d5 100644
--- 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;
 }