enable jit compilation caching (still testing)

programs/e3make3d_gauss.py

@@ -41,6 +41,15 @@
 # used for profiling. This can be commented out as long as the @profile line is also commented out
 #from line_profiler import profile
 
+try: os.mkdir(".jaxcache")
+except: pass
+
+# We cache the JIT compilation results to speed up future runs
+jax.config.update("jax_compilation_cache_dir", "./.jaxcache")
+jax.config.update("jax_persistent_cache_min_entry_size_bytes", -1)
+jax.config.update("jax_persistent_cache_min_compile_time_secs", 2)
+jax.config.update("jax_persistent_cache_enable_xla_caches", "xla_gpu_per_fusion_autotune_cache_dir")
+
 jax.config.update("jax_default_matmul_precision", "float32")
 
 # @profile