enviroment variable JAX fast time check

Author: Jammy2211

autofit/jax_wrapper.py

@@ -1,17 +1,42 @@
+import logging
+
+logger = logging.getLogger(__name__)
+
 """
 Allows the user to switch between using NumPy and JAX for linear algebra operations.
 
 If USE_JAX=true in general.yaml then JAX's NumPy is used, otherwise vanilla NumPy is used.
 """
-import jax
-import os
-
 from autoconf import conf
 
 use_jax = conf.instance["general"]["jax"]["use_jax"]
 
 if use_jax:
 
+    import os
+
+    xla_env = os.environ.get("XLA_FLAGS")
+
+    if "--xla_disable_hlo_passes=constant_folding" not in xla_env:
+        logger.info(
+            """
+            For fast JAX compile times, the envirment variable XLA_FLAGS must be set to "--xla_disable_hlo_passes=constant_folding",
+            which is currently not.
+            
+            In Python, to do this manually, use the code: 
+            
+            import os
+            os.environ["XLA_FLAGS"] = "--xla_disable_hlo_passes=constant_folding"
+            
+            The environment variable has been set automatically for you now, however if JAX has already been imported, 
+            this change will not take effect and JAX function compiling times may be slow. 
+            
+            Therefore, it is recommended to set this environment variable before running your script, e.g. in your terminal.
+            """)
+
+        os.environ['XLA_FLAGS'] = "--xla_disable_hlo_passes=constant_folding"
+
+    import jax
     from jax import numpy
 
     print(