Skip to content

Commit

Permalink
still trying to fix docs
Browse files Browse the repository at this point in the history
  • Loading branch information
YigitElma committed Dec 20, 2024
1 parent 8737f7f commit 601da8a
Show file tree
Hide file tree
Showing 2 changed files with 3 additions and 2 deletions.
2 changes: 1 addition & 1 deletion docs/index.rst
Original file line number Diff line number Diff line change
Expand Up @@ -21,6 +21,7 @@
:caption: Tutorials

installation
performance_tips
notebooks/tutorials/basic_equilibrium.ipynb
notebooks/tutorials/advanced_equilibrium_continuation.ipynb
notebooks/tutorials/continuation_step_by_step.ipynb
Expand All @@ -35,7 +36,6 @@
notebooks/tutorials/coil_stage_two_optimization.ipynb
notebooks/tutorials/QFM_surface.ipynb
notebooks/tutorials/ideal_ballooning_stability.ipynb
performance_tips

.. toctree::
:maxdepth: 1
Expand Down
3 changes: 2 additions & 1 deletion docs/performance_tips.rst
Original file line number Diff line number Diff line change
Expand Up @@ -7,14 +7,15 @@ Caching the Compiled (Jitted) Code
Although the compiled code is fast, it still takes time to compile. If you are running the same optimization, or some similar optimization, multiple times, you can save time by caching the compiled code. This automatically happens for a single session (for example, until you restart your kernel in Jupyter Notebook) but once you start using another session, the code will need to be recompiled. Fortunately, there is a way to bypass this. First create a cache directory (i.e. ``jax-caches``), and put the following code at the beginning of your script:

.. code-block:: python
import jax
import jax.numpy as jnp
jax.config.update("jax_compilation_cache_dir", "../jax-caches")
jax.config.update("jax_persistent_cache_min_entry_size_bytes", -1)
jax.config.update("jax_persistent_cache_min_compile_time_secs", 0)
This will use a directory called ``jax-caches`` in the parent directory of the script to store the compiled code. The ``jax_persistent_cache_min_entry_size_bytes`` and ``jax_persistent_cache_min_compile_time_secs`` parameters are set to -1 and 0, respectively, to ensure that all compiled code is cached. For more details on caching, refer to official JAX documentation `here <https://jax.readthedocs.io/en/latest/persistent_compilation_cache.html#persistent-compilation-cache>`__.
This will use a directory called ``jax-caches`` in the parent directory of the script to store the compiled code. The ``jax_persistent_cache_min_entry_size_bytes`` and ``jax_persistent_cache_min_compile_time_secs`` parameters are set to -1 and 0, respectively, to ensure that all compiled code is cached. For more details on caching, refer to official JAX documentation [`here <https://jax.readthedocs.io/en/latest/persistent_compilation_cache.html#persistent-compilation-cache>`__].

Note: Updating JAX version might re-compile some previously cached code, and thi might increase the cache size. Every once in a while, you might need to clear your cache directory.

Expand Down

0 comments on commit 601da8a

Please sign in to comment.