TST: Speed up repeated jax tests (#649)

Add pytest option --jax-pcc for the jax persistent compilation cache
and add notes to developer contributing guide
This uses a jax config to set cache options.
Total size of cache for all tests (well, really four main ones) is only
2.7M on linux, 416 files.
Reduces runtime from
4.16s call     test_optimizers.py::test_fit[SBR]
2.98s call     test_optimizers.py::test_complexity_not_fitted[SBR]
2.89s call     test_optimizers.py::test_pickle[SBR-opt_args1]
2.89s call     test_optimizers.py::test_sbr_accurate

To:
1.94s call     test_optimizers.py::test_fit[SBR]
1.56s call     test_optimizers.py::test_sbr_accurate
1.23s call     test_optimizers.py::test_pickle[SBR-opt_args1]
1.18s call     test_optimizers.py::test_complexity_not_fitted[SBR]

Author: Jacob-Stevens-Haas

docs/contributing.rst

@@ -141,15 +141,16 @@ Install our pre-commit script via
 
     pre-commit install
 
+Pre-commit will automatically check all future commits for code style.
 To be accepted your code should conform to PEP8 and pass all unit tests.
 Code can be tested by invoking
 
 .. code-block:: bash
 
-    pytest
-
-Pre-commit will automatically check all future commits for code style.
+    pytest --jax-pcc
 
+The ``jax-pcc`` flag is optional and will speed up the jax tests on repeated runs
+by caching compiled functions to disk.
 
 
 Coding Guidelines

test/conftest.py

@@ -2,6 +2,7 @@
 Shared pytest fixtures for unit tests.
 """
 from pathlib import Path
+from tempfile import gettempdir
 
 import jax
 import numpy as np
@@ -34,6 +35,14 @@ def pytest_addoption(parser):
             " test_notebooks.test_external"
         ),
     )
+    parser.addoption(
+        "--jax-pcc",
+        action="store_true",
+        help=(
+            "Whether to cache @jax.jit compilations to disk."
+            "It speeds up the slowest tests appx 2x, writing files to /tmp/jax_cache"
+        ),
+    )
 
 
 def pytest_generate_tests(metafunc):
@@ -47,6 +56,17 @@ def pytest_generate_tests(metafunc):
         )
 
 
+@pytest.fixture(scope="session")
+def set_jax_pcc_env(request):
+    if request.config.getoption("--jax-pcc"):
+        cache_dir = Path(gettempdir()) / "jax_cache/"
+        cache_dir.mkdir(exist_ok=True)
+        jax.config.update("jax_compilation_cache_dir", str(cache_dir))
+        jax.config.update("jax_persistent_cache_min_entry_size_bytes", -1)
+        jax.config.update("jax_persistent_cache_min_compile_time_secs", -1)
+        jax.config.update("jax_persistent_cache_enable_xla_caches", "all")
+
+
 @pytest.fixture(scope="session")
 def data_1d():
     t = np.linspace(0, 1, 10)

test/test_optimizers/test_optimizers.py

@@ -130,7 +130,7 @@ def data(request):
     ],
     ids=lambda param: type(param),
 )
-def test_fit(data_derivative_1d, optimizer):
+def test_fit(data_derivative_1d, optimizer, set_jax_pcc_env):
     x, x_dot = data_derivative_1d
     if len(x.shape) == 1:
         x = x.reshape(-1, 1)
@@ -168,7 +168,7 @@ def test_not_fitted(optimizer):
     ],
     ids=type,
 )
-def test_complexity_not_fitted(optimizer, data_derivative_2d):
+def test_complexity_not_fitted(optimizer, data_derivative_2d, set_jax_pcc_env):
     with pytest.raises(NotFittedError):
         optimizer.complexity
 
@@ -389,7 +389,7 @@ def test_sbr_bad_parameters(params):
         SBR(**params)
 
 
-def test_sbr_accurate():
+def test_sbr_accurate(set_jax_pcc_env):
     # It's really hard to tune SBR to get desired shrinkage
     # This just tests that SBR fits "close" to unregularized regression
     x = np.tile(np.eye(2), 4).reshape((-1, 2))
@@ -1193,7 +1193,7 @@ def test_remove_and_decrement():
         (TrappingSR3, {"_n_tgts": 3, "_include_bias": True}),
     ),
 )
-def test_pickle(data_lorenz, opt_cls, opt_args):
+def test_pickle(data_lorenz, opt_cls, opt_args, set_jax_pcc_env):
     x, t = data_lorenz
     y = PolynomialLibrary(degree=2).fit_transform(x)
     opt = opt_cls(**opt_args).fit(y, x)