diff --git a/README.md b/README.md index 7278733..8ce7ad2 100644 --- a/README.md +++ b/README.md @@ -1,151 +1,188 @@ Hyperoptax Logo -# Hyperoptax: Parallel hyperparameter tuning with JAX +# Hyperoptax: parallel hyperparameter tuning with JAX [![PyPI version](https://img.shields.io/pypi/v/hyperoptax)](https://pypi.org/project/hyperoptax) ![CI status](https://github.com/TheodoreWolf/hyperoptax/actions/workflows/test.yml/badge.svg?branch=main) [![codecov](https://codecov.io/gh/TheodoreWolf/hyperoptax/graph/badge.svg?token=Y582MZ25GG)](https://codecov.io/gh/TheodoreWolf/hyperoptax) +Hyperoptax is a small, functional hyperparameter-optimization library for JAX. +It supports nested PyTree search spaces, batched candidate evaluation, Python and +`jax.lax.scan` loops, and random, grid, and Gaussian-process Bayesian search. -## ⛰️ Introduction - -Hyperoptax is a lightweight toolbox for parallel hyperparameter optimization of pure JAX functions. It provides a concise API that lets you wrap any JAX-compatible loss or evaluation function and search across spaces __in parallel__ – all while staying in pure JAX. - -## 🏗️ Installation +## Installation ```bash -pip install hyperoptax +uv pip install hyperoptax ``` -If you want to use the notebooks: -```bash -pip install hyperoptax[notebooks] -``` - -If you do not yet have JAX installed, pick the right wheel for your accelerator: +Install the optional notebook dependencies with: ```bash -# CPU-only -pip install --upgrade "jax[cpu]" -# or GPU/TPU – see the official JAX installation guide +uv pip install "hyperoptax[notebooks]" ``` -## 🥜 In a nutshell - +Hyperoptax requires Python 3.10 or newer. For accelerator-specific JAX wheels, +follow the [JAX installation guide](https://docs.jax.dev/en/latest/installation.html). -All optimizers follow the same stateless pattern: `Optimizer.init` returns a `(state, optimizer)` pair, and `optimizer.optimize` runs the search loop. Your objective function must have the signature `fn(key, params) -> scalar`. Importantly, `params` can be _any_ PyTree. +## Quick start +Every optimizer follows the same functional pattern: `init` returns an immutable +state PyTree and an optimizer, and the objective has signature +`objective(key, params) -> scalar`. ```python import jax -from hyperoptax import BayesianSearch, LogSpace, LinearSpace +import jax.numpy as jnp + +from hyperoptax import BayesianSearch, LinearSpace, LogSpace -def train_nn(key, params): - learning_rate = params["learning_rate"] - final_lr_pct = params["final_lr_pct"] - ... - return val_loss # scalar, lower is better -search_space = { +def objective(key, params): + noise = 0.01 * jax.random.normal(key) + return ( + (jnp.log10(params["learning_rate"]) + 3.0) ** 2 + + (params["dropout"] - 0.2) ** 2 + + noise + ) + + +space = { "learning_rate": LogSpace(1e-5, 1e-1), - "final_lr_pct": LinearSpace(0.01, 0.5), + "dropout": LinearSpace(0.0, 0.5), } state, optimizer = BayesianSearch.init( - search_space, - n_max=100, # observation buffer size (= number of iterations) - n_parallel=4, # Parallel workers per step + space, + n_max=80, + n_parallel=4, maximize=False, ) - -state, (params_hist, results_hist) = optimizer.optimize( - state, jax.random.PRNGKey(0), train_nn +state, (params_history, results_history) = optimizer.optimize_scan( + state, + jax.random.PRNGKey(0), + objective, + n_iterations=20, ) -# params_hist: list of pytrees, one per iteration (each leaf has shape (n_parallel,)) -# results_hist: list of arrays, one per iteration (each has shape (n_parallel,)) -# Retrieve best result print(optimizer.best_result(state)) print(optimizer.best_params(state)) ``` -Other available optimizers: +Here `n_iterations=20` and `n_parallel=4` perform 80 objective evaluations. +For buffered optimizers, `n_max` is the number of observations the state can +store, not the number of iterations. Keep `n_max` divisible by `n_parallel` to +avoid evaluating a partially storable final batch. + +## Choosing an optimizer + +| Optimizer | Best suited to | Important constraints | +|---|---|---| +| `RandomSearch` | Cheap baseline and broad exploration | Stateless; does not retain observations | +| `GridSearch` | Small, explicitly enumerated grids | Every leaf must be a numeric `DiscreteSpace`; make the grid size divisible by `n_parallel` | +| `BayesianSearch` | Expensive, smooth objectives | Fixed `n_max` buffer; GP cost grows with the buffer size | + +All spaces may be nested in arbitrary dict/list/tuple PyTrees. Available leaf +types are `LinearSpace`, `LogSpace`, `DiscreteSpace`, `QLinearSpace`, and +`QLogSpace`. `DiscreteSpace` values must be numeric because optimizers represent +candidates as JAX arrays. + +### Random and grid search ```python -from hyperoptax import RandomSearch, GridSearch, DiscreteSpace +from hyperoptax import DiscreteSpace, GridSearch, RandomSearch -# Random search -state, optimizer = RandomSearch.init(search_space, n_parallel=8) -state, history = optimizer.optimize(state, jax.random.PRNGKey(0), train_nn, n_iterations=50) +# Random search: 10 batches of 8 independent samples. +random_state, random_search = RandomSearch.init(space, n_parallel=8) +random_state, random_history = random_search.optimize( + random_state, jax.random.PRNGKey(1), objective, n_iterations=10 +) -# Grid search (DiscreteSpace only) -# Note: shuffle=True -grid_space = {"lr": DiscreteSpace([1e-4, 1e-3, 1e-2]), "dropout": DiscreteSpace([0.1, 0.3, 0.5])} -state, optimizer = GridSearch.init(grid_space) -state, history = optimizer.optimize(state, jax.random.PRNGKey(0), train_nn, n_iterations=9) +# Grid search: six points, evaluated two at a time. +grid_space = { + "width": DiscreteSpace((64, 128, 256)), + "depth": DiscreteSpace((2, 4)), +} +grid_state, grid_search = GridSearch.init(grid_space, n_parallel=2) ``` -### `optimize_scan()` — JAX-native loop - -`optimize_scan()` has the same signature as `optimize()` but uses `jax.lax.scan` internally. -This requires your objective function to be JAX-traceable (jit-compilable), and returns -**stacked arrays** rather than Python lists: +Pass `shuffle=True` and an explicit PRNG key to randomize grid traversal +reproducibly: ```python -state, (params_hist, results_hist) = optimizer.optimize_scan( - state, jax.random.PRNGKey(0), train_nn, n_iterations=25 +grid_state, grid_search = GridSearch.init( + grid_space, + key=jax.random.PRNGKey(2), + shuffle=True, + n_parallel=2, ) -# params_hist: pytree where each leaf has shape (n_iterations, n_parallel, ...) -# results_hist: array of shape (n_iterations, n_parallel) ``` -> **Return type difference:** `optimize()` returns Python lists (easy to index by iteration), -> while `optimize_scan()` returns stacked JAX arrays (compatible with `jax.jit`, faster for -> JAX-traceable objectives). Choose based on your objective function and use case. +## Python loop versus JAX scan + +`optimize()` runs the outer loop in Python and returns lists indexed by +iteration: + +- `params_history`: list of PyTrees whose leaves have leading shape + `(n_parallel, ...)`. +- `results_history`: list of arrays with shape `(n_parallel,)`. -## 💪 Hyperoptax in action -BayesOpt animation +`optimize_scan()` uses `jax.lax.scan`, requires a JAX-traceable objective, and +returns stacked arrays: -## 🔪 The Sharp Bits +- `params_history`: PyTree with leaf shape + `(n_iterations, n_parallel, ...)`. +- `results_history`: array with shape `(n_iterations, n_parallel)`. -Since we are working in pure JAX the same [sharp bits](https://docs.jax.dev/en/latest/notebooks/Common_Gotchas_in_JAX.html) apply. Some consequences of this for hyperoptax: -1. Parameters that change the length of an evaluation (e.g: epochs, generations...) can't be optimized in parallel. -2. Neural network structures can't be optimized in parallel either. -3. Strings can't be used as hyperparameters. +Both paths use `jax.vmap` to evaluate each batch. Use the lower-level +`get_next_params` and `update_state` methods when evaluations must run in an +external system. -## 🫂 Contributing +## JAX constraints -We welcome pull requests! To get started: +The usual [JAX sharp bits](https://docs.jax.dev/en/latest/notebooks/Common_Gotchas_in_JAX.html) +apply: -1. Open an issue describing the bug or feature. -2. Fork the repository and create a feature branch (`git checkout -b user/my-feature`). -3. Clone and install dependencies. We recommend [uv](https://docs.astral.sh/uv/) for environment management: +1. The objective must accept a PRNG key and one unbatched parameter PyTree, and + return a scalar. +2. Every batch is evaluated with `jax.vmap`; data-dependent Python control flow + and side effects are not supported. +3. Strings and variable-shaped model structures cannot be represented as + batched JAX hyperparameters. Encode categorical choices numerically and map + them to host-side labels where needed. +4. Parameters that change evaluation shape or duration can interact poorly with + parallel batches; use `n_parallel=1` when evaluations cannot be synchronized. + +## Notebooks + +See [`notebooks/`](notebooks/) for grid/Bayesian search examples, design +studies, high-dimensional behavior, performance comparisons, RL tuning, and GP +visualization. + +## Contributing ```bash git clone https://github.com/TheodoreWolf/hyperoptax cd hyperoptax uv pip install -e ".[all]" +uv run pytest -m "not timing" +uv run ruff check src tests +uv run ruff format --check src tests ``` -4. Run the test suite: - -```bash -uv run pytest -``` -5. Ensure the notebooks still work. -6. Format your code with `ruff`. -7. Submit a pull request. +The full suite includes optional timing tests; run `uv run pytest` when you want +those as well. Please keep objectives reproducible by passing and splitting JAX +PRNG keys explicitly. ## Roadmap -I'm developing this both as a passion project and for my work in my PhD. I have a few ideas on where to go with this library: -- Callbacks! -- Reduce redundant kernel recomputation — currently the full K matrix is rebuilt each iteration when only the new row/column is needed. -- Length scale tuning currently uses a fixed Adam step count; smarter convergence criteria could help. -- Tree Parzen Estimator (TPE), this is essentially SOTA for hyperparameter search, implementing this would be super cool! -## 📝 Citation +- Add callbacks for logging and early stopping. +- Reuse GP kernel blocks instead of rebuilding the full matrix each iteration. +- Replace the fixed number of length-scale Adam steps with a convergence-aware + stopping criterion. +- Add more optimizer families for mixed and cost-sensitive search spaces. -If you use Hyperoptax in academic work, please cite: +## Citation ```bibtex @misc{hyperoptax, diff --git a/docs/requirements.txt b/docs/requirements.txt index 9f32705..ed0bbd2 100644 --- a/docs/requirements.txt +++ b/docs/requirements.txt @@ -1,5 +1,3 @@ sphinx>=6.1 -myst-parser -nbsphinx sphinx-book-theme>=1.1.4 -sphinx-autobuild \ No newline at end of file +sphinx-autobuild diff --git a/docs/source/api/_autosummary/hyperoptax.acquisition.rst b/docs/source/api/_autosummary/hyperoptax.acquisition.rst deleted file mode 100644 index 2fee26d..0000000 --- a/docs/source/api/_autosummary/hyperoptax.acquisition.rst +++ /dev/null @@ -1,14 +0,0 @@ -hyperoptax.acquisition -====================== - -.. automodule:: hyperoptax.acquisition - - - .. rubric:: Classes - - .. autosummary:: - - BaseAcquisition - EI - UCB - \ No newline at end of file diff --git a/docs/source/api/_autosummary/hyperoptax.base.rst b/docs/source/api/_autosummary/hyperoptax.base.rst deleted file mode 100644 index 22dab21..0000000 --- a/docs/source/api/_autosummary/hyperoptax.base.rst +++ /dev/null @@ -1,12 +0,0 @@ -hyperoptax.base -=============== - -.. automodule:: hyperoptax.base - - - .. rubric:: Classes - - .. autosummary:: - - BaseOptimizer - \ No newline at end of file diff --git a/docs/source/api/_autosummary/hyperoptax.bayesian.rst b/docs/source/api/_autosummary/hyperoptax.bayesian.rst deleted file mode 100644 index 010af4f..0000000 --- a/docs/source/api/_autosummary/hyperoptax.bayesian.rst +++ /dev/null @@ -1,12 +0,0 @@ -hyperoptax.bayesian -=================== - -.. automodule:: hyperoptax.bayesian - - - .. rubric:: Classes - - .. autosummary:: - - BayesianOptimizer - \ No newline at end of file diff --git a/docs/source/api/_autosummary/hyperoptax.grid.rst b/docs/source/api/_autosummary/hyperoptax.grid.rst deleted file mode 100644 index 7e7e6c9..0000000 --- a/docs/source/api/_autosummary/hyperoptax.grid.rst +++ /dev/null @@ -1,12 +0,0 @@ -hyperoptax.grid -=============== - -.. automodule:: hyperoptax.grid - - - .. rubric:: Classes - - .. autosummary:: - - GridSearch - \ No newline at end of file diff --git a/docs/source/api/_autosummary/hyperoptax.kernels.rst b/docs/source/api/_autosummary/hyperoptax.kernels.rst deleted file mode 100644 index 5d4aaad..0000000 --- a/docs/source/api/_autosummary/hyperoptax.kernels.rst +++ /dev/null @@ -1,20 +0,0 @@ -hyperoptax.kernels -================== - -.. automodule:: hyperoptax.kernels - - - .. rubric:: Functions - - .. autosummary:: - - cdist - - .. rubric:: Classes - - .. autosummary:: - - BaseKernel - Matern - RBF - \ No newline at end of file diff --git a/docs/source/api/_autosummary/hyperoptax.spaces.rst b/docs/source/api/_autosummary/hyperoptax.spaces.rst deleted file mode 100644 index a690764..0000000 --- a/docs/source/api/_autosummary/hyperoptax.spaces.rst +++ /dev/null @@ -1,17 +0,0 @@ -hyperoptax.spaces -================= - -.. automodule:: hyperoptax.spaces - - - .. rubric:: Classes - - .. autosummary:: - - ArbitrarySpace - BaseSpace - ExpSpace - LinearSpace - LogSpace - QuantizedLinearSpace - \ No newline at end of file diff --git a/docs/source/api/acquisition.rst b/docs/source/api/acquisition.rst index 58a169a..439903c 100644 --- a/docs/source/api/acquisition.rst +++ b/docs/source/api/acquisition.rst @@ -1,32 +1,37 @@ -Acquisition Functions -==================== +Acquisition and hallucination strategies +======================================== -Acquisition functions determine which points to evaluate next in Bayesian optimization. +Bayesian search scores candidate points with an acquisition function. +Probability of improvement is the default; expected improvement and upper +confidence bound are also available. -.. automodule:: hyperoptax.acquisition - :members: - :undoc-members: - :show-inheritance: - -Usage Examples --------------- - -Upper Confidence Bound (UCB) -~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ +Parallel Bayesian batches are selected sequentially. After each selection, a +hallucination strategy supplies a temporary objective value so the following +slot explores a different part of the posterior. Posterior sampling is the +default; mean, UCB, and constant strategies are available. .. code-block:: python - from hyperoptax.acquisition import UCB - - # Create UCB acquisition function - acq = UCB(kappa=2.0) + from hyperoptax import BayesianSearch, EI, MeanHallucination, UCB -Expected Improvement (EI) -~~~~~~~~~~~~~~~~~~~~~~~~~ + state, optimizer = BayesianSearch.init( + space, + n_max=80, + n_parallel=4, + acquisition=EI(xi=0.01), + hallucination=MeanHallucination(), + ) -.. code-block:: python + # UCB is an alternative acquisition strategy. + state, optimizer = BayesianSearch.init( + space, + n_max=80, + acquisition=UCB(kappa=2.0), + ) + +API +--- - from hyperoptax.acquisition import EI - - # Create EI acquisition function - acq = EI(xi=0.01) \ No newline at end of file +.. automodule:: hyperoptax.acquisition + :members: + :show-inheritance: diff --git a/docs/source/api/hyperoptax.rst b/docs/source/api/hyperoptax.rst deleted file mode 100644 index 1869796..0000000 --- a/docs/source/api/hyperoptax.rst +++ /dev/null @@ -1,58 +0,0 @@ -hyperoptax package -================== - -.. automodule:: hyperoptax - :members: - :undoc-members: - :show-inheritance: - -Submodules ----------- - -hyperoptax.acquisition module ------------------------------ - -.. automodule:: hyperoptax.acquisition - :members: - :undoc-members: - :show-inheritance: - -hyperoptax.base module ----------------------- - -.. automodule:: hyperoptax.base - :members: - :undoc-members: - :show-inheritance: - -hyperoptax.bayesian module --------------------------- - -.. automodule:: hyperoptax.bayesian - :members: - :undoc-members: - :show-inheritance: - -hyperoptax.grid module ----------------------- - -.. automodule:: hyperoptax.grid - :members: - :undoc-members: - :show-inheritance: - -hyperoptax.kernels module -------------------------- - -.. automodule:: hyperoptax.kernels - :members: - :undoc-members: - :show-inheritance: - -hyperoptax.spaces module ------------------------- - -.. automodule:: hyperoptax.spaces - :members: - :undoc-members: - :show-inheritance: diff --git a/docs/source/api/kernels.rst b/docs/source/api/kernels.rst index 9958fba..9ba60f8 100644 --- a/docs/source/api/kernels.rst +++ b/docs/source/api/kernels.rst @@ -1,32 +1,28 @@ Kernels ======= -Kernels are used by the Gaussian process in Bayesian optimization to model function similarity. +Bayesian search uses a Matérn kernel with ``nu=0.5`` by default. RBF and Matérn +kernels accept scalar initial length scales; Bayesian search maintains and tunes +per-dimension ARD length scales in optimizer state. -.. automodule:: hyperoptax.kernels - :members: - :undoc-members: - :show-inheritance: - -Usage Examples --------------- +.. code-block:: python -RBF Kernel -~~~~~~~~~~~ + from hyperoptax import BayesianSearch, Matern, RBF -.. code-block:: python + state, optimizer = BayesianSearch.init( + space, + n_max=80, + kernel=Matern(length_scale=1.0, nu=1.5), + ) - from hyperoptax.kernels import RBF - - # Create RBF kernel with length scale 1.0 - kernel = RBF(length_scale=1.0) + rbf = RBF(length_scale=0.5) -Matern Kernel -~~~~~~~~~~~~~ +The supported Matérn smoothness values are ``0.5``, ``1.5``, ``2.5``, and +``float("inf")``. The infinite-smoothness case is equivalent to RBF. -.. code-block:: python +API +--- - from hyperoptax.kernels import Matern - - # Create Matern kernel with custom parameters - kernel = Matern(length_scale=1.0, nu=2.5) \ No newline at end of file +.. automodule:: hyperoptax.kernels + :members: + :show-inheritance: diff --git a/docs/source/api/modules.rst b/docs/source/api/modules.rst deleted file mode 100644 index fc28c8c..0000000 --- a/docs/source/api/modules.rst +++ /dev/null @@ -1,7 +0,0 @@ -hyperoptax -========== - -.. toctree:: - :maxdepth: 4 - - hyperoptax diff --git a/docs/source/api/optimizers.rst b/docs/source/api/optimizers.rst index dc4e0fb..92aebd5 100644 --- a/docs/source/api/optimizers.rst +++ b/docs/source/api/optimizers.rst @@ -1,36 +1,34 @@ Optimizers ========== -Hyperoptax provides several optimization algorithms for hyperparameter tuning. +All optimizers implement the functional interface defined by +:class:`hyperoptax.base.Optimizer`: initialize state, request a batch, evaluate +it, then update state with the results. -Base Optimizer +Base interface -------------- .. automodule:: hyperoptax.base :members: - :undoc-members: :show-inheritance: -Bayesian Optimizer ------------------- +Random search +------------- -.. automodule:: hyperoptax.bayesian +.. automodule:: hyperoptax.random :members: - :undoc-members: :show-inheritance: -Grid Search +Grid search ----------- .. automodule:: hyperoptax.grid :members: - :undoc-members: :show-inheritance: -Random Search -------------- +Bayesian search +--------------- -.. automodule:: hyperoptax.random +.. automodule:: hyperoptax.bayesian :members: - :undoc-members: - :show-inheritance: \ No newline at end of file + :show-inheritance: diff --git a/docs/source/api/spaces.rst b/docs/source/api/spaces.rst index 32c76a2..8bde3a1 100644 --- a/docs/source/api/spaces.rst +++ b/docs/source/api/spaces.rst @@ -1,40 +1,39 @@ -Parameter Spaces -================ +Search spaces +============= -Parameter spaces define the search domains for hyperparameter optimization. - -.. automodule:: hyperoptax.spaces - :members: - :undoc-members: - :show-inheritance: - -Examples --------- - -Creating a Linear Space -~~~~~~~~~~~~~~~~~~~~~~~ +Search spaces are JAX PyTree leaves and may be nested in dicts, lists, tuples, +or other PyTree containers. Candidate batches preserve the same structure and +add a leading ``n_parallel`` axis to every leaf. .. code-block:: python - from hyperoptax import LinearSpace + import jax.numpy as jnp - dropout_space = LinearSpace(0.0, 0.5) + from hyperoptax import ( + DiscreteSpace, + LinearSpace, + LogSpace, + QLinearSpace, + QLogSpace, + ) -Creating a Logarithmic Space -~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ - -.. code-block:: python + space = { + "optimizer": { + "learning_rate": LogSpace(1e-5, 1e-1), + "weight_decay": LinearSpace(0.0, 0.1), + }, + "depth": QLinearSpace(2, 8, datatype=jnp.int32), + "width": QLogSpace(32, 1024, datatype=jnp.int32), + "batch_size": DiscreteSpace((32, 64, 128, 256)), + } - from hyperoptax import LogSpace +``DiscreteSpace`` values must be numeric. Encode string categories as numbers +and map them back to labels outside the vmapped objective when needed. - lr_space = LogSpace(1e-5, 1e-1) +API +--- -Creating a Discrete Space -~~~~~~~~~~~~~~~~~~~~~~~~~~ - -.. code-block:: python - - from hyperoptax import DiscreteSpace - - optimizer_space = DiscreteSpace(["adam", "sgd", "rmsprop"]) - lr_space = DiscreteSpace([1e-4, 1e-3, 1e-2]) \ No newline at end of file +.. automodule:: hyperoptax.spaces + :members: + :exclude-members: datatype + :show-inheritance: diff --git a/docs/source/conf.py b/docs/source/conf.py index 82863e1..9d82e6b 100644 --- a/docs/source/conf.py +++ b/docs/source/conf.py @@ -5,6 +5,7 @@ import os import sys +from importlib.metadata import version as package_version # Add the source directory to Python path so Sphinx can find the modules sys.path.insert(0, os.path.abspath("../../src")) @@ -13,38 +14,30 @@ # https://www.sphinx-doc.org/en/master/usage/configuration.html#project-information project = "Hyperoptax" -copyright = "2025, Theo Wolf" +copyright = "2025-2026, Theo Wolf" author = "Theo Wolf" -release = "0.1.6" +release = package_version("hyperoptax") +version = release # -- General configuration --------------------------------------------------- # https://www.sphinx-doc.org/en/master/usage/configuration.html#general-configuration extensions = [ - "sphinx.ext.autodoc", # Auto-generate docs from docstrings - "sphinx.ext.autosummary", # Generate summary tables - "sphinx.ext.viewcode", # Add source code links - "sphinx.ext.napoleon", # Support for NumPy/Google style docstrings - "sphinx.ext.intersphinx", # Link to other project docs - # "myst_parser", # Markdown support (temporarily disabled) - # "nbsphinx", # Jupyter notebook support (may cause issues) + "sphinx.ext.autodoc", + "sphinx.ext.viewcode", + "sphinx.ext.napoleon", + "sphinx.ext.intersphinx", ] -templates_path = ["_templates"] -exclude_patterns = [] +exclude_patterns = ["_build"] # Autodoc configuration autodoc_default_options = { "members": True, "member-order": "bysource", - "special-members": "__init__", - "undoc-members": True, - "exclude-members": "__weakref__", + "show-inheritance": True, } -# Autosummary configuration -autosummary_generate = True - # Include type hints in the description rather than the signature so that # Sphinx Napoleon + autodoc produce cleaner function/class signatures. autodoc_typehints = "description" @@ -58,7 +51,7 @@ # Intersphinx mapping to link to external docs intersphinx_mapping = { "python": ("https://docs.python.org/3/", None), - "jax": ("https://jax.readthedocs.io/en/latest/", None), + "jax": ("https://docs.jax.dev/en/latest/", None), "numpy": ("https://numpy.org/doc/stable/", None), } @@ -84,8 +77,8 @@ # Disable expansion to keep a fixed navigation layout "collapse_navbar": True, "logo": { - "image_light": "_static/manifold.png", - "image_dark": "_static/manifold.png", + "image_light": "../../assets/logo-transparent.png", + "image_dark": "../../assets/logo-transparent.png", "text": "Hyperoptax", "alt_text": "Hyperoptax - Parallel hyperparameter tuning with JAX", }, diff --git a/docs/source/index.rst b/docs/source/index.rst index 81bfd46..d170e5e 100644 --- a/docs/source/index.rst +++ b/docs/source/index.rst @@ -1,56 +1,69 @@ -.. Hyperoptax documentation master file, created by - sphinx-quickstart on Tue Jul 15 22:39:09 2025. - You can adapt this file completely to your liking, but it should at least - contain the root `toctree` directive. +Hyperoptax +========== -.. raw:: html +Hyperoptax is a functional hyperparameter-optimization library for JAX. It +supports PyTree search spaces, parallel candidate evaluation, Python and +``jax.lax.scan`` loops, and random, grid, and Gaussian-process Bayesian search. -
- 🚧 WORK IN PROGRESS - This documentation is currently under development 🚧 -
+Quick start +----------- -Hyperoptax Documentation -======================== +Every objective has signature ``objective(key, params) -> scalar``. Optimizers +are initialized as an immutable state PyTree plus a stateless optimizer object. -Welcome to Hyperoptax - a lightweight toolbox for parallel hyperparameter optimization of pure JAX functions. +.. code-block:: python -Hyperoptax provides a concise API that lets you wrap any JAX-compatible loss or evaluation function and search across parameter spaces **in parallel** - all while staying in pure JAX. + import jax + import jax.numpy as jnp -Quick Start ------------ + from hyperoptax import BayesianSearch, LinearSpace, LogSpace -.. code-block:: python - import jax - from hyperoptax import BayesianSearch, LogSpace, LinearSpace + def objective(key, params): + noise = 0.01 * jax.random.normal(key) + return ( + (jnp.log10(params["learning_rate"]) + 3.0) ** 2 + + (params["dropout"] - 0.2) ** 2 + + noise + ) - def train_nn(key, params): - learning_rate = params["learning_rate"] - final_lr_pct = params["final_lr_pct"] - ... - return val_loss # scalar, lower is better - search_space = { + space = { "learning_rate": LogSpace(1e-5, 1e-1), - "final_lr_pct": LinearSpace(0.01, 0.5), + "dropout": LinearSpace(0.0, 0.5), } - state, optimizer = BayesianSearch.init( - search_space, - n_max=100, + space, + n_max=80, + n_parallel=4, maximize=False, ) - state, (params_hist, results_hist) = optimizer.optimize( - state, jax.random.PRNGKey(0), train_nn + state, history = optimizer.optimize_scan( + state, + jax.random.PRNGKey(0), + objective, + n_iterations=20, ) + + print(optimizer.best_result(state)) print(optimizer.best_params(state)) +The example performs 80 evaluations: 20 iterations with four candidates per +iteration. See :doc:`usage` for buffer semantics, optimizer selection, and the +lower-level ask/tell-style API. + +.. toctree:: + :maxdepth: 2 + :caption: User guide + + usage + notebooks/index + .. toctree:: :maxdepth: 2 - :titlesonly: + :caption: API reference api/optimizers api/spaces - api/kernels api/acquisition - + api/kernels diff --git a/docs/source/notebooks/index.rst b/docs/source/notebooks/index.rst index 9a03473..d4c48b1 100644 --- a/docs/source/notebooks/index.rst +++ b/docs/source/notebooks/index.rst @@ -1,41 +1,24 @@ -Examples and Tutorials -====================== - -Learn how to use Hyperoptax through practical examples and tutorials. - -.. toctree:: - :maxdepth: 1 - :caption: Examples: - - ../../../notebooks/search - ../../../notebooks/performance - ../../../notebooks/rl_hparams - ../../../notebooks/visualization - -Quick Examples --------------- - -Grid Search Example -~~~~~~~~~~~~~~~~~~~ - -Basic usage of grid search for hyperparameter optimization. - -Bayesian Optimization Example -~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ - -Using Bayesian optimization with Gaussian processes for efficient hyperparameter search. - -Performance Comparison -~~~~~~~~~~~~~~~~~~~~~ - -Comparing Hyperoptax with other optimization libraries. - -Reinforcement Learning -~~~~~~~~~~~~~~~~~~~~~ - -Hyperparameter tuning for RL algorithms using Rejax. - -Visualization -~~~~~~~~~~~~ - -Visualizing the Bayesian optimization process. \ No newline at end of file +Notebooks +========= + +The repository contains executable studies rather than embedding notebook +outputs in the documentation build: + +* `Search examples `_: + grid search and Bayesian optimization. +* `Bayesian optimization design study `_: + acquisition functions and their hyperparameters. +* `High-dimensional behavior `_: + Bayesian versus random search from 10 to 1,000 dimensions. +* `Performance comparison `_: + comparisons with other optimization libraries and batch sizes. +* `RL hyperparameters `_: + tuning Rejax reinforcement-learning runs. +* `Bayesian optimization visualization `_: + GP posterior and acquisition behavior. + +Install the optional dependencies before running them: + +.. code-block:: console + + uv pip install -e ".[notebooks]" diff --git a/docs/source/usage.rst b/docs/source/usage.rst new file mode 100644 index 0000000..01d8498 --- /dev/null +++ b/docs/source/usage.rst @@ -0,0 +1,103 @@ +Using Hyperoptax +================ + +Installation +------------ + +.. code-block:: console + + uv pip install hyperoptax + +Hyperoptax requires Python 3.10 or newer. Install accelerator-specific JAX +wheels by following the `JAX installation guide +`_. + +Choosing an optimizer +--------------------- + +.. list-table:: + :header-rows: 1 + :widths: 18 32 50 + + * - Optimizer + - Best suited to + - Constraints + * - :class:`~hyperoptax.random.RandomSearch` + - Cheap baseline and broad exploration + - Stateless; does not retain observations. + * - :class:`~hyperoptax.grid.GridSearch` + - Small, explicitly enumerated grids + - Every leaf must be a numeric + :class:`~hyperoptax.spaces.DiscreteSpace`. The grid size should be + divisible by ``n_parallel``. + * - :class:`~hyperoptax.bayesian.BayesianSearch` + - Expensive, smooth objectives + - Stores observations in a fixed ``n_max`` buffer; GP work grows with the + buffer size. + +Iterations, batches, and buffers +-------------------------------- + +``n_parallel`` is the number of configurations evaluated in each optimizer +iteration. Therefore ``n_iterations * n_parallel`` is the number of objective +evaluations. + +Bayesian search uses fixed-size buffers so its state has static JAX shapes. +``n_max`` is the number of observations that fit in the buffer, not the number +of iterations. When ``n_iterations`` is omitted, Bayesian search runs until the +buffer is full. Keep ``n_max`` divisible by ``n_parallel`` so the last batch +fits exactly. + +Grid search builds the full Cartesian product during ``init``. If the grid size +is not divisible by ``n_parallel``, its incomplete final batch is omitted. + +Loop APIs +--------- + +:meth:`~hyperoptax.base.Optimizer.optimize` runs the outer loop in Python. It +returns ``params_history`` and ``results_history`` as lists indexed by +iteration. The parameter PyTrees have a leading ``n_parallel`` axis and each +result array has shape ``(n_parallel,)``. + +:meth:`~hyperoptax.base.Optimizer.optimize_scan` uses ``jax.lax.scan`` and +requires a JAX-traceable objective. It returns stacked parameter leaves with +shape ``(n_iterations, n_parallel, ...)`` and results with shape +``(n_iterations, n_parallel)``. + +Both APIs evaluate the batch with ``jax.vmap``. The objective must accept one +unbatched configuration and return a scalar. + +External evaluation loop +------------------------ + +Use :meth:`~hyperoptax.base.Optimizer.get_next_params` and +:meth:`~hyperoptax.base.Optimizer.update_state` directly when a job scheduler, +service, or other external system evaluates configurations. + +.. code-block:: python + + import jax + + key = jax.random.PRNGKey(0) + previous_params = None + previous_results = None + + for _ in range(20): + key, ask_key, eval_key, tell_key = jax.random.split(key, 4) + params = optimizer.get_next_params( + state, ask_key, previous_params, previous_results + ) + eval_keys = jax.random.split(eval_key, optimizer.n_parallel) + results = jax.vmap(objective)(eval_keys, params) + state = optimizer.update_state(state, tell_key, results, params) + previous_params, previous_results = params, results + +JAX constraints +--------------- + +* Search-space leaves and objective results must be representable as JAX + arrays. String categories are not supported; encode them numerically. +* Data-dependent Python control flow, variable-shaped outputs, and side effects + are not compatible with ``optimize_scan``. +* Parameters that change evaluation shape or duration may not batch cleanly. + Use ``n_parallel=1`` when evaluations cannot be synchronized. diff --git a/pyproject.toml b/pyproject.toml index ae6303a..6d5eea9 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -48,8 +48,6 @@ build-backend = "setuptools.build_meta" [dependency-groups] docs = [ "sphinx>=6.1", - "myst-parser", - "nbsphinx", "sphinx-book-theme>=1.1.4", "sphinx-autobuild" ]