From bf89615c9f01707a116cc7741d5121775e3dc856 Mon Sep 17 00:00:00 2001 From: Theo Wolf Date: Sat, 11 Apr 2026 14:03:59 +0100 Subject: [PATCH 1/3] Remove flax dependancy --- pyproject.toml | 4 ++-- src/hyperoptax/base.py | 21 +++++++++---------- src/hyperoptax/bayesian.py | 29 ++++++++++++++------------- src/hyperoptax/grid.py | 9 +++++---- src/hyperoptax/random.py | 4 ++-- src/hyperoptax/spaces.py | 41 ++++++++++++++++++++++++++------------ src/hyperoptax/utils.py | 7 ++++--- tests/test_bayes.py | 40 ++++++++++++++++++------------------- tests/test_grid.py | 4 +++- tests/test_spaces.py | 7 +++---- 10 files changed, 93 insertions(+), 73 deletions(-) diff --git a/pyproject.toml b/pyproject.toml index 2818b9a..5db5956 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -4,7 +4,7 @@ version = "0.1.6" description = "Tuning hyperparameters with JAX" readme = "README.md" requires-python = ">=3.10" -dependencies = ["jax>=0.4.38", "flax", "optax"] +dependencies = ["jax>=0.4.38", "optax"] license = "Apache-2.0" authors = [{name = "Theo Wolf"}] @@ -49,7 +49,7 @@ build-backend = "setuptools.build_meta" docs = [ "sphinx>=6.1", "myst-parser", - "nbsphinx", + "nbsphinx", "sphinx-book-theme>=1.1.4", "sphinx-autobuild" ] diff --git a/src/hyperoptax/base.py b/src/hyperoptax/base.py index 91d3300..b38c5a7 100644 --- a/src/hyperoptax/base.py +++ b/src/hyperoptax/base.py @@ -1,10 +1,10 @@ import inspect import warnings -from typing import Callable +from dataclasses import dataclass +from typing import Any, Callable import jax import jax.numpy as jnp -from flax import struct def _validate_func(func): @@ -35,11 +35,12 @@ def _validate_func(func): return -@struct.dataclass +@jax.tree_util.register_dataclass +@dataclass(frozen=True) class OptimizerState: - """Base optimizer state — a Flax PyTree holding the search space definition.""" + """Base optimizer state — a JAX pytree holding the search space definition.""" - space: struct.PyTreeNode + space: Any class Optimizer: @@ -55,7 +56,7 @@ def optimize( key: jax.Array, # () PRNG key func: Callable, # (key, config) -> () scalar result n_iterations: int, - ) -> tuple[OptimizerState, tuple[struct.PyTreeNode, jax.Array]]: + ) -> tuple[OptimizerState, tuple[Any, jax.Array]]: """ High Level API for optimizing a function over a space. Not recommended if you want to do fancy things @@ -86,7 +87,7 @@ def optimize_scan( key: jax.Array, # () PRNG key func: Callable, # (key, config) -> () scalar result n_iterations: int, - ) -> tuple[OptimizerState, tuple[struct.PyTreeNode, jax.Array]]: + ) -> tuple[OptimizerState, tuple[Any, jax.Array]]: """ Like optimize, but uses jax.lax.scan for the inner loop. @@ -140,7 +141,7 @@ def update_state( state: OptimizerState, key: jax.random.PRNGKey, results: jax.Array, - params: struct.PyTreeNode = None, + params: Any = None, ) -> OptimizerState: """ Updates the optimizer state based on the results of the function. @@ -151,9 +152,9 @@ def get_next_params( self, state: OptimizerState, key: jax.random.PRNGKey, - params: struct.PyTreeNode = None, + params: Any = None, results: jax.Array = None, - ) -> struct.PyTreeNode: + ) -> Any: """ Gets the next parameters to sample from the space. Returns a batched pytree where every leaf has shape (n_parallel, ...). diff --git a/src/hyperoptax/bayesian.py b/src/hyperoptax/bayesian.py index b192f7b..06db889 100644 --- a/src/hyperoptax/bayesian.py +++ b/src/hyperoptax/bayesian.py @@ -4,7 +4,6 @@ import jax import jax.numpy as jnp import optax -from flax import struct from hyperoptax import acquisition as acq from hyperoptax import base, kernels @@ -13,7 +12,8 @@ MASK_VARIANCE = 1e12 # large diagonal added to masked rows to isolate them from GP fit -@struct.dataclass +@jax.tree_util.register_dataclass +@dataclasses.dataclass(frozen=True) class BayesianSearchState(base.OptimizerState): """State for :class:`BayesianSearch`. @@ -41,9 +41,9 @@ class BayesianSearchState(base.OptimizerState): class BayesianSearch(base.Optimizer): """Bayesian optimisation with a Gaussian Process surrogate. - Uses a GP (Matérn 2.5 kernel by default) to model the objective and + Uses a GP (Matérn 0.5 kernel by default) to model the objective and selects the next batch of candidates by maximising an acquisition function - (EI by default). ARD length scales are tuned with Adam each iteration. + (PI by default). ARD length scales are tuned with Adam each iteration. Parallel batches are generated via the Kriging Believer hallucination strategy. @@ -51,9 +51,9 @@ class BayesianSearch(base.Optimizer): jitter: Small diagonal added to the kernel matrix for numerical stability (default ``1e-6``). kernel: Kernel function (default :class:`~hyperoptax.kernels.Matern` - with ``nu=2.5``). + with ``nu=0.5``). acquisition: Acquisition function (default - :class:`~hyperoptax.acquisition.EI` with ``xi=0.01``). + :class:`~hyperoptax.acquisition.PI` with ``xi=0.01``). n_candidates: Number of random candidates sampled per iteration for the discrete pre-selection step (default ``1000``). n_restarts: Number of L-BFGS restarts seeded from the top candidates @@ -64,17 +64,17 @@ class BayesianSearch(base.Optimizer): n_warmup: Number of pure-random iterations before the GP is used (default ``1``). maximize: Set ``False`` to minimise the objective (default ``True``). - n_parallel: Number of parallel candidates per iteration (default ``1``). + n_parallel: Number of parallel candidates per iteration (default ``4``). hallucination: Hallucination strategy for Kriging Believer parallel - selection (default :class:`~hyperoptax.acquisition.MeanHallucination`). + selection (default :class:`~hyperoptax.acquisition.SampleHallucination`). """ jitter: float = 1e-6 kernel: kernels.BaseKernel = dataclasses.field( - default_factory=lambda: kernels.Matern(length_scale=1.0, nu=2.5) + default_factory=lambda: kernels.Matern(length_scale=1.0, nu=0.5) ) acquisition: acq.BaseAcquisition = dataclasses.field( - default_factory=lambda: acq.EI(xi=0.01) + default_factory=lambda: acq.PI(xi=0.01) ) n_candidates: int = 1000 # random candidates sampled for continuous spaces n_restarts: int = 2 # number of L-BFGS restarts (seeded from top candidates) @@ -82,9 +82,9 @@ class BayesianSearch(base.Optimizer): n_hparam_steps: int = 20 # Adam steps to tune log_length_scale each iteration n_warmup: int = 1 # pure-random evaluations before GP kicks in maximize: bool = True # set False to minimize the objective - n_parallel: int = 1 + n_parallel: int = 4 hallucination: acq.BaseHallucination = dataclasses.field( - default_factory=acq.MeanHallucination + default_factory=acq.SampleHallucination ) @classmethod @@ -382,7 +382,8 @@ def body(i, s): y_scalar = jax.lax.dynamic_slice(results, (i,), (1,)) return jax.lax.cond( slot < n_max, - lambda s: s.replace( + lambda s: dataclasses.replace( + s, X=jax.lax.dynamic_update_slice(s.X, x_row, (slot, 0)), y=jax.lax.dynamic_update_slice(s.y, y_scalar, (slot,)), mask=jax.lax.dynamic_update_slice( @@ -428,7 +429,7 @@ def update_state(self, state, key, results, params): lambda s: s.log_length_scale, state, ) - state = state.replace(log_length_scale=log_ls) + state = dataclasses.replace(state, log_length_scale=log_ls) return state def _n_iterations(self, state): diff --git a/src/hyperoptax/grid.py b/src/hyperoptax/grid.py index 17e62f2..56bd99f 100644 --- a/src/hyperoptax/grid.py +++ b/src/hyperoptax/grid.py @@ -1,14 +1,15 @@ import dataclasses +from typing import Any import jax import jax.numpy as jnp -from flax import struct from hyperoptax import base from hyperoptax import spaces as sp -@struct.dataclass +@jax.tree_util.register_dataclass +@dataclasses.dataclass(frozen=True) class GridSearchState(base.OptimizerState): """State for :class:`GridSearch`. @@ -86,7 +87,7 @@ def init(cls, space, key=None, **kwargs): def get_next_params( self, state: GridSearchState, key, params=None, results=None - ) -> struct.PyTreeNode: + ) -> Any: """Return the next ``n_parallel`` parameter combinations from the grid.""" # Only check eagerly; inside lax.scan grid_idx is an abstract tracer. if not isinstance(state.grid_idx, jax.core.Tracer): @@ -110,4 +111,4 @@ def update_state( self, state: GridSearchState, key, results, params=None ) -> GridSearchState: """Advance the grid index by ``n_parallel``.""" - return state.replace(grid_idx=state.grid_idx + self.n_parallel) + return dataclasses.replace(state, grid_idx=state.grid_idx + self.n_parallel) diff --git a/src/hyperoptax/random.py b/src/hyperoptax/random.py index 9345bc2..8ed9b4f 100644 --- a/src/hyperoptax/random.py +++ b/src/hyperoptax/random.py @@ -1,8 +1,8 @@ import dataclasses +from typing import Any import jax import jax.numpy as jnp -from flax import struct from hyperoptax import base, utils from hyperoptax import spaces as sp @@ -31,7 +31,7 @@ def get_next_params( key: jax.random.PRNGKey, params=None, results=None, - ) -> struct.PyTreeNode: + ) -> Any: """Sample ``n_parallel`` independent configurations from the search space.""" def sample_once(k): subkeys = utils.make_key_tree(state.space, k) diff --git a/src/hyperoptax/spaces.py b/src/hyperoptax/spaces.py index df49eac..2341177 100644 --- a/src/hyperoptax/spaces.py +++ b/src/hyperoptax/spaces.py @@ -1,6 +1,8 @@ +from abc import ABC, abstractmethod +from dataclasses import dataclass, field + import jax import jax.numpy as jnp -from flax import struct # transformation between logs @@ -8,9 +10,11 @@ def log_transform(x: float, base: float) -> float: return jnp.log(x) / jnp.log(base) -class Space(struct.PyTreeNode): +@dataclass(frozen=True) +class Space(ABC): """Abstract base class for hyperparameter search spaces.""" + @abstractmethod def sample(self, key: jax.random.PRNGKey) -> jax.Array: raise NotImplementedError @@ -18,6 +22,8 @@ def transform(self, value): return value +@jax.tree_util.register_dataclass +@dataclass(frozen=True) class LinearSpace(Space): """Uniform continuous space over ``[lower_bound, upper_bound]``. @@ -26,8 +32,8 @@ class LinearSpace(Space): upper_bound: Exclusive upper bound of the interval. """ - lower_bound: float = struct.field(pytree_node=False) - upper_bound: float = struct.field(pytree_node=False) + lower_bound: float = field(metadata=dict(static=True)) + upper_bound: float = field(metadata=dict(static=True)) def __post_init__(self): assert self.lower_bound < self.upper_bound, ( @@ -42,6 +48,8 @@ def sample(self, key: jax.random.PRNGKey) -> float: ) +@jax.tree_util.register_dataclass +@dataclass(frozen=True) class DiscreteSpace(Space): """Discrete space over a fixed set of values. @@ -53,7 +61,7 @@ class DiscreteSpace(Space): values: Tuple of candidate values to sample from. """ - values: tuple = struct.field(pytree_node=False) + values: tuple = field(metadata=dict(static=True)) @property def lower_bound(self) -> float: @@ -76,6 +84,8 @@ def transform(self, value) -> jax.Array: return snapped.reshape(value.shape) +@jax.tree_util.register_dataclass +@dataclass(frozen=True) class LogSpace(LinearSpace): """Log-uniform continuous space over ``[lower_bound, upper_bound]``. @@ -89,15 +99,13 @@ class LogSpace(LinearSpace): base: Logarithm base (default ``10``). Must be greater than 1. """ - base: float = struct.field(pytree_node=False, default=10) + base: float = field(default=10, metadata=dict(static=True)) - def __post_init__( - self, - ): + def __post_init__(self): super().__post_init__() assert self.base > 1, "Log base must be greater than 1" - def sample(self, key: jax.random.PRNGKey) -> float: + def sample(self, key: jax.random.PRNGKey) -> jax.Array: return self.transform( self.base ** jax.random.uniform( @@ -111,6 +119,8 @@ def sample(self, key: jax.random.PRNGKey) -> float: # TODO: maybe use something more robust than astype? # TODO: can we do something with mixins? Currently hitting some ordering problems +@jax.tree_util.register_dataclass +@dataclass(frozen=True) class QLinearSpace(LinearSpace): """Quantized (integer) variant of :class:`LinearSpace`. @@ -124,13 +134,15 @@ class QLinearSpace(LinearSpace): datatype: Integer dtype used after rounding (default ``jnp.int32``). """ - datatype: type = struct.field(pytree_node=False, default=jnp.int32) + datatype: type = field(default=jnp.int32, metadata=dict(static=True)) def transform(self, value) -> jax.Array: return jnp.round(value).astype(self.datatype) -class QLogSpace(LogSpace, QLinearSpace): +@jax.tree_util.register_dataclass +@dataclass(frozen=True) +class QLogSpace(LogSpace): """Quantized (integer) variant of :class:`LogSpace`. Samples in log space and rounds to the nearest integer. Use this for @@ -138,4 +150,7 @@ class QLogSpace(LogSpace, QLinearSpace): (e.g. number of hidden units, number of warmup steps). """ - pass + datatype: type = field(default=jnp.int32, metadata=dict(static=True)) + + def transform(self, value) -> jax.Array: + return jnp.round(value).astype(self.datatype) diff --git a/src/hyperoptax/utils.py b/src/hyperoptax/utils.py index 30b551c..77438fd 100644 --- a/src/hyperoptax/utils.py +++ b/src/hyperoptax/utils.py @@ -1,13 +1,14 @@ +from typing import Any + import jax -from flax import struct from hyperoptax import spaces as sp def make_key_tree( - pytree: struct.PyTreeNode, + pytree: Any, subkey: jax.random.PRNGKey, -) -> struct.PyTreeNode: +) -> Any: """Split ``subkey`` into a pytree of PRNGKeys matching the structure of ``pytree``. :class:`~hyperoptax.spaces.Space` objects are treated as leaves, so the diff --git a/tests/test_bayes.py b/tests/test_bayes.py index 92797d5..3290072 100644 --- a/tests/test_bayes.py +++ b/tests/test_bayes.py @@ -315,8 +315,8 @@ def test_2d_space(self): params = optimizer.get_next_params(state, self.key) assert "x" in params assert "y" in params - assert params["x"].shape == (1,) - assert params["y"].shape == (1,) + assert params["x"].shape == (4,) + assert params["y"].shape == (4,) def test_n_parallel_discrete(self): space = {"x": sp.DiscreteSpace([0.0, 0.25, 0.5, 0.75, 1.0])} @@ -328,7 +328,7 @@ def test_n_parallel_discrete(self): class TestBayesianSearchOptimize: def test_optimize_returns_correct_shapes(self): space = {"x": sp.DiscreteSpace([0.0, 0.25, 0.5, 0.75, 1.0])} - state, optimizer = bayesian.BayesianSearch.init(space, n_max=5) + state, optimizer = bayesian.BayesianSearch.init(space, n_max=5, n_parallel=1) func = lambda key, config: -(config["x"] ** 2) state, (params_hist, results_hist) = optimizer.optimize( state, jax.random.PRNGKey(0), func @@ -338,14 +338,14 @@ def test_optimize_returns_correct_shapes(self): def test_optimize_fills_state(self): space = {"x": sp.DiscreteSpace([0.0, 0.25, 0.5, 0.75, 1.0])} - state, optimizer = bayesian.BayesianSearch.init(space, n_max=5) + state, optimizer = bayesian.BayesianSearch.init(space, n_max=5, n_parallel=1) func = lambda key, config: -(config["x"] ** 2) state, _ = optimizer.optimize(state, jax.random.PRNGKey(0), func) assert int(state.mask.sum()) == 5 def test_optimize_finds_optimum(self): space = {"x": sp.DiscreteSpace([0.0, 0.25, 0.5, 0.75, 1.0])} - state, optimizer = bayesian.BayesianSearch.init(space, n_max=20) + state, optimizer = bayesian.BayesianSearch.init(space, n_max=20, n_parallel=1) func = lambda key, config: -(config["x"] ** 2) state, (params_hist, results_hist) = optimizer.optimize( state, jax.random.PRNGKey(0), func @@ -355,7 +355,7 @@ def test_optimize_finds_optimum(self): def test_optimize_with_array_result(self): space = {"x": sp.DiscreteSpace([0.0, 0.25, 0.5, 0.75, 1.0])} - state, optimizer = bayesian.BayesianSearch.init(space, n_max=3) + state, optimizer = bayesian.BayesianSearch.init(space, n_max=3, n_parallel=1) func = lambda key, config: jnp.array([-(config["x"] ** 2)]) state, (_, results_hist) = optimizer.optimize( state, jax.random.PRNGKey(0), func @@ -365,7 +365,7 @@ def test_optimize_with_array_result(self): def test_optimize_converges_toward_optimum(self): space = {"x": sp.DiscreteSpace([0.0, 0.25, 0.5, 0.75, 1.0])} state, optimizer = bayesian.BayesianSearch.init( - space, n_max=20, acquisition=acq.UCB() + space, n_max=20, n_parallel=1, acquisition=acq.UCB() ) func = lambda key, config: -(config["x"] ** 2) state, _ = optimizer.optimize(state, jax.random.PRNGKey(0), func) @@ -373,7 +373,7 @@ def test_optimize_converges_toward_optimum(self): def test_optimize_continuous_space(self): space = {"x": sp.LinearSpace(0.0, 1.0)} - state, optimizer = bayesian.BayesianSearch.init(space, n_max=5) + state, optimizer = bayesian.BayesianSearch.init(space, n_max=5, n_parallel=1) func = lambda key, config: -(config["x"] ** 2) state, (params_hist, results_hist) = optimizer.optimize( state, jax.random.PRNGKey(0), func @@ -384,7 +384,7 @@ def test_optimize_continuous_space(self): def test_optimize_continuous_with_ei_uses_observed_y_max(self): space = {"x": sp.LinearSpace(0.0, 1.0), "y": sp.LinearSpace(0.0, 1.0)} state, optimizer = bayesian.BayesianSearch.init( - space, n_max=20, acquisition=acq.EI() + space, n_max=20, n_parallel=1, acquisition=acq.EI() ) state = optimizer.update_state( state, jax.random.PRNGKey(0), jnp.array([100.0]), jnp.array([[0.5, 0.5]]) @@ -394,10 +394,10 @@ def test_optimize_continuous_with_ei_uses_observed_y_max(self): def test_optimize_minimize(self): space = {"x": sp.DiscreteSpace([0.0, 0.25, 0.5, 0.75, 1.0])} - state, optimizer = bayesian.BayesianSearch.init(space, n_max=10, maximize=False) + state, optimizer = bayesian.BayesianSearch.init(space, n_max=20, n_parallel=1, maximize=False) func = lambda key, config: config["x"] ** 2 state, _ = optimizer.optimize(state, jax.random.PRNGKey(0), func) - assert int(state.mask.sum()) == 10 + assert int(state.mask.sum()) == 20 assert float(optimizer.best_result(state)) == pytest.approx(0.0) def test_optimize_n_parallel_fills_buffer(self): @@ -458,7 +458,7 @@ def test_best_params_minimize(self): assert float(params["x"]) == pytest.approx(0.5) def test_best_result_after_full_optimize(self): - state, optimizer = bayesian.BayesianSearch.init(self.space, n_max=10) + state, optimizer = bayesian.BayesianSearch.init(self.space, n_max=10, n_parallel=1) func = lambda key, config: -(config["x"] ** 2) state, _ = optimizer.optimize(state, jax.random.PRNGKey(0), func) assert float(optimizer.best_result(state)) == pytest.approx( @@ -509,7 +509,7 @@ def test_log_length_scale_unchanged_with_single_observation(self): def test_tuned_length_scale_used_in_gp(self): state, optimizer = bayesian.BayesianSearch.init( - self.space, n_max=10, n_hparam_steps=20 + self.space, n_max=10, n_parallel=1, n_hparam_steps=20 ) func = lambda key, config: -(config["x"] ** 2) state, _ = optimizer.optimize(state, jax.random.PRNGKey(0), func) @@ -551,7 +551,7 @@ def test_get_next_params_mixed(self): "lr": sp.LogSpace(1e-4, 1e-1), "layers": sp.DiscreteSpace([1, 2, 3, 4]), } - state, optimizer = bayesian.BayesianSearch.init(space, n_max=20) + state, optimizer = bayesian.BayesianSearch.init(space, n_max=20, n_parallel=1) params = optimizer.get_next_params(state, jax.random.PRNGKey(0)) assert "lr" in params and "layers" in params assert params["lr"].shape == (1,) @@ -593,7 +593,7 @@ def test_two_arg_passes(self): class TestBayesianSearchOptimizeScan: def test_optimize_scan_runs(self): space = {"x": sp.LinearSpace(0.0, 1.0)} - state, optimizer = bayesian.BayesianSearch.init(space, n_max=5) + state, optimizer = bayesian.BayesianSearch.init(space, n_max=5, n_parallel=1) func = lambda key, config: -(config["x"] ** 2) state, (params_hist, results_hist) = optimizer.optimize_scan( state, jax.random.PRNGKey(0), func @@ -719,9 +719,9 @@ def test_lbfgs_improves_over_seed(self): state, key, jnp.array([r]), jnp.array([[x, y_]]) ) params = optimizer.get_next_params(state, key) - assert params["x"].shape == (1,) - assert 0.0 <= float(params["x"][0]) <= 1.0 - assert 0.0 <= float(params["y"][0]) <= 1.0 + assert params["x"].shape == (optimizer.n_parallel,) + assert all(0.0 <= float(v) <= 1.0 for v in params["x"]) + assert all(0.0 <= float(v) <= 1.0 for v in params["y"]) class TestKrigingBelieverHallucination: @@ -749,10 +749,10 @@ def _make_state_with_obs( state = optimizer.update_state(state, key, jnp.array([float(i)]), x) return state, optimizer - def test_default_hallucination_is_mean(self): + def test_default_hallucination_is_sample(self): space = {"x": sp.LinearSpace(0.0, 1.0)} _, optimizer = bayesian.BayesianSearch.init(space) - assert isinstance(optimizer.hallucination, acq.MeanHallucination) + assert isinstance(optimizer.hallucination, acq.SampleHallucination) def test_mean_hallucination_optimize_runs(self): space = {"x": sp.LinearSpace(0.0, 1.0)} diff --git a/tests/test_grid.py b/tests/test_grid.py index a949df5..532adc5 100644 --- a/tests/test_grid.py +++ b/tests/test_grid.py @@ -1,3 +1,5 @@ +import dataclasses + import jax import jax.numpy as jnp import pytest @@ -157,7 +159,7 @@ def test_get_next_params_raises_on_overflow(self): space = {"x": sp.DiscreteSpace([0.0, 0.5, 1.0])} state, optimizer = grid.GridSearch.init(space, n_parallel=2) # Move to last valid position - state = state.replace(grid_idx=2) + state = dataclasses.replace(state, grid_idx=2) with pytest.raises(ValueError, match="Not enough grid points"): optimizer.get_next_params(state, _KEY) diff --git a/tests/test_spaces.py b/tests/test_spaces.py index e144851..07851ba 100644 --- a/tests/test_spaces.py +++ b/tests/test_spaces.py @@ -95,11 +95,10 @@ def test_space_post_init(): sp.LogSpace(0, 10, base=1) -def test_space_base_sample_raises(): +def test_space_base_cannot_instantiate(): # Space.sample is abstract — subclasses must override it - space = sp.Space() - with pytest.raises(NotImplementedError): - space.sample(jax.random.PRNGKey(0)) + with pytest.raises(TypeError): + sp.Space() def test_discrete_space_lower_upper_bound(): From 954a7bbe4add36e2e27ac6afe4b2e5a1f149f479 Mon Sep 17 00:00:00 2001 From: Theo Wolf Date: Sat, 11 Apr 2026 14:55:37 +0100 Subject: [PATCH 2/3] Add jaxtyping as dependancy --- pyproject.toml | 2 +- src/hyperoptax/base.py | 19 ++++++++++++------- src/hyperoptax/bayesian.py | 5 ++--- src/hyperoptax/grid.py | 6 +++--- src/hyperoptax/random.py | 4 ++-- src/hyperoptax/utils.py | 7 +++---- tests/test_grid.py | 3 +-- 7 files changed, 24 insertions(+), 22 deletions(-) diff --git a/pyproject.toml b/pyproject.toml index 5db5956..759ff5d 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -4,7 +4,7 @@ version = "0.1.6" description = "Tuning hyperparameters with JAX" readme = "README.md" requires-python = ">=3.10" -dependencies = ["jax>=0.4.38", "optax"] +dependencies = ["jax>=0.4.38", "optax", "jaxtyping"] license = "Apache-2.0" authors = [{name = "Theo Wolf"}] diff --git a/src/hyperoptax/base.py b/src/hyperoptax/base.py index b38c5a7..e455e6b 100644 --- a/src/hyperoptax/base.py +++ b/src/hyperoptax/base.py @@ -1,10 +1,12 @@ +import dataclasses import inspect import warnings from dataclasses import dataclass -from typing import Any, Callable +from typing import Callable import jax import jax.numpy as jnp +from jaxtyping import PyTree def _validate_func(func): @@ -40,7 +42,10 @@ def _validate_func(func): class OptimizerState: """Base optimizer state — a JAX pytree holding the search space definition.""" - space: Any + space: PyTree + + def replace(self, **kwargs) -> "OptimizerState": + return dataclasses.replace(self, **kwargs) class Optimizer: @@ -56,7 +61,7 @@ def optimize( key: jax.Array, # () PRNG key func: Callable, # (key, config) -> () scalar result n_iterations: int, - ) -> tuple[OptimizerState, tuple[Any, jax.Array]]: + ) -> tuple[OptimizerState, tuple[PyTree, jax.Array]]: """ High Level API for optimizing a function over a space. Not recommended if you want to do fancy things @@ -87,7 +92,7 @@ def optimize_scan( key: jax.Array, # () PRNG key func: Callable, # (key, config) -> () scalar result n_iterations: int, - ) -> tuple[OptimizerState, tuple[Any, jax.Array]]: + ) -> tuple[OptimizerState, tuple[PyTree, jax.Array]]: """ Like optimize, but uses jax.lax.scan for the inner loop. @@ -141,7 +146,7 @@ def update_state( state: OptimizerState, key: jax.random.PRNGKey, results: jax.Array, - params: Any = None, + params: PyTree | None = None, ) -> OptimizerState: """ Updates the optimizer state based on the results of the function. @@ -152,9 +157,9 @@ def get_next_params( self, state: OptimizerState, key: jax.random.PRNGKey, - params: Any = None, + params: PyTree | None = None, results: jax.Array = None, - ) -> Any: + ) -> PyTree: """ Gets the next parameters to sample from the space. Returns a batched pytree where every leaf has shape (n_parallel, ...). diff --git a/src/hyperoptax/bayesian.py b/src/hyperoptax/bayesian.py index 06db889..5f14706 100644 --- a/src/hyperoptax/bayesian.py +++ b/src/hyperoptax/bayesian.py @@ -382,8 +382,7 @@ def body(i, s): y_scalar = jax.lax.dynamic_slice(results, (i,), (1,)) return jax.lax.cond( slot < n_max, - lambda s: dataclasses.replace( - s, + lambda s: s.replace( X=jax.lax.dynamic_update_slice(s.X, x_row, (slot, 0)), y=jax.lax.dynamic_update_slice(s.y, y_scalar, (slot,)), mask=jax.lax.dynamic_update_slice( @@ -429,7 +428,7 @@ def update_state(self, state, key, results, params): lambda s: s.log_length_scale, state, ) - state = dataclasses.replace(state, log_length_scale=log_ls) + state = state.replace(log_length_scale=log_ls) return state def _n_iterations(self, state): diff --git a/src/hyperoptax/grid.py b/src/hyperoptax/grid.py index 56bd99f..401e6b3 100644 --- a/src/hyperoptax/grid.py +++ b/src/hyperoptax/grid.py @@ -1,8 +1,8 @@ import dataclasses -from typing import Any import jax import jax.numpy as jnp +from jaxtyping import PyTree from hyperoptax import base from hyperoptax import spaces as sp @@ -87,7 +87,7 @@ def init(cls, space, key=None, **kwargs): def get_next_params( self, state: GridSearchState, key, params=None, results=None - ) -> Any: + ) -> PyTree: """Return the next ``n_parallel`` parameter combinations from the grid.""" # Only check eagerly; inside lax.scan grid_idx is an abstract tracer. if not isinstance(state.grid_idx, jax.core.Tracer): @@ -111,4 +111,4 @@ def update_state( self, state: GridSearchState, key, results, params=None ) -> GridSearchState: """Advance the grid index by ``n_parallel``.""" - return dataclasses.replace(state, grid_idx=state.grid_idx + self.n_parallel) + return state.replace(grid_idx=state.grid_idx + self.n_parallel) diff --git a/src/hyperoptax/random.py b/src/hyperoptax/random.py index 8ed9b4f..af4964e 100644 --- a/src/hyperoptax/random.py +++ b/src/hyperoptax/random.py @@ -1,8 +1,8 @@ import dataclasses -from typing import Any import jax import jax.numpy as jnp +from jaxtyping import PyTree from hyperoptax import base, utils from hyperoptax import spaces as sp @@ -31,7 +31,7 @@ def get_next_params( key: jax.random.PRNGKey, params=None, results=None, - ) -> Any: + ) -> PyTree: """Sample ``n_parallel`` independent configurations from the search space.""" def sample_once(k): subkeys = utils.make_key_tree(state.space, k) diff --git a/src/hyperoptax/utils.py b/src/hyperoptax/utils.py index 77438fd..efa6b30 100644 --- a/src/hyperoptax/utils.py +++ b/src/hyperoptax/utils.py @@ -1,14 +1,13 @@ -from typing import Any - import jax +from jaxtyping import PyTree from hyperoptax import spaces as sp def make_key_tree( - pytree: Any, + pytree: PyTree, subkey: jax.random.PRNGKey, -) -> Any: +) -> PyTree: """Split ``subkey`` into a pytree of PRNGKeys matching the structure of ``pytree``. :class:`~hyperoptax.spaces.Space` objects are treated as leaves, so the diff --git a/tests/test_grid.py b/tests/test_grid.py index 532adc5..03a0729 100644 --- a/tests/test_grid.py +++ b/tests/test_grid.py @@ -1,4 +1,3 @@ -import dataclasses import jax import jax.numpy as jnp @@ -159,7 +158,7 @@ def test_get_next_params_raises_on_overflow(self): space = {"x": sp.DiscreteSpace([0.0, 0.5, 1.0])} state, optimizer = grid.GridSearch.init(space, n_parallel=2) # Move to last valid position - state = dataclasses.replace(state, grid_idx=2) + state = state.replace(grid_idx=2) with pytest.raises(ValueError, match="Not enough grid points"): optimizer.get_next_params(state, _KEY) From 959e4afd9dcd796bd1e847fd03e38a6f63a5e1a1 Mon Sep 17 00:00:00 2001 From: Theo Wolf Date: Sat, 11 Apr 2026 15:09:35 +0100 Subject: [PATCH 3/3] Clean up branch --- src/hyperoptax/bayesian.py | 20 +++++++++---------- tests/test_bayes.py | 40 +++++++++++++++++++------------------- 2 files changed, 30 insertions(+), 30 deletions(-) diff --git a/src/hyperoptax/bayesian.py b/src/hyperoptax/bayesian.py index 5f14706..e9f20e5 100644 --- a/src/hyperoptax/bayesian.py +++ b/src/hyperoptax/bayesian.py @@ -41,9 +41,9 @@ class BayesianSearchState(base.OptimizerState): class BayesianSearch(base.Optimizer): """Bayesian optimisation with a Gaussian Process surrogate. - Uses a GP (Matérn 0.5 kernel by default) to model the objective and + Uses a GP (Matérn 2.5 kernel by default) to model the objective and selects the next batch of candidates by maximising an acquisition function - (PI by default). ARD length scales are tuned with Adam each iteration. + (EI by default). ARD length scales are tuned with Adam each iteration. Parallel batches are generated via the Kriging Believer hallucination strategy. @@ -51,9 +51,9 @@ class BayesianSearch(base.Optimizer): jitter: Small diagonal added to the kernel matrix for numerical stability (default ``1e-6``). kernel: Kernel function (default :class:`~hyperoptax.kernels.Matern` - with ``nu=0.5``). + with ``nu=2.5``). acquisition: Acquisition function (default - :class:`~hyperoptax.acquisition.PI` with ``xi=0.01``). + :class:`~hyperoptax.acquisition.EI` with ``xi=0.01``). n_candidates: Number of random candidates sampled per iteration for the discrete pre-selection step (default ``1000``). n_restarts: Number of L-BFGS restarts seeded from the top candidates @@ -64,17 +64,17 @@ class BayesianSearch(base.Optimizer): n_warmup: Number of pure-random iterations before the GP is used (default ``1``). maximize: Set ``False`` to minimise the objective (default ``True``). - n_parallel: Number of parallel candidates per iteration (default ``4``). + n_parallel: Number of parallel candidates per iteration (default ``1``). hallucination: Hallucination strategy for Kriging Believer parallel - selection (default :class:`~hyperoptax.acquisition.SampleHallucination`). + selection (default :class:`~hyperoptax.acquisition.MeanHallucination`). """ jitter: float = 1e-6 kernel: kernels.BaseKernel = dataclasses.field( - default_factory=lambda: kernels.Matern(length_scale=1.0, nu=0.5) + default_factory=lambda: kernels.Matern(length_scale=1.0, nu=2.5) ) acquisition: acq.BaseAcquisition = dataclasses.field( - default_factory=lambda: acq.PI(xi=0.01) + default_factory=lambda: acq.EI(xi=0.01) ) n_candidates: int = 1000 # random candidates sampled for continuous spaces n_restarts: int = 2 # number of L-BFGS restarts (seeded from top candidates) @@ -82,9 +82,9 @@ class BayesianSearch(base.Optimizer): n_hparam_steps: int = 20 # Adam steps to tune log_length_scale each iteration n_warmup: int = 1 # pure-random evaluations before GP kicks in maximize: bool = True # set False to minimize the objective - n_parallel: int = 4 + n_parallel: int = 1 hallucination: acq.BaseHallucination = dataclasses.field( - default_factory=acq.SampleHallucination + default_factory=acq.MeanHallucination ) @classmethod diff --git a/tests/test_bayes.py b/tests/test_bayes.py index 3290072..92797d5 100644 --- a/tests/test_bayes.py +++ b/tests/test_bayes.py @@ -315,8 +315,8 @@ def test_2d_space(self): params = optimizer.get_next_params(state, self.key) assert "x" in params assert "y" in params - assert params["x"].shape == (4,) - assert params["y"].shape == (4,) + assert params["x"].shape == (1,) + assert params["y"].shape == (1,) def test_n_parallel_discrete(self): space = {"x": sp.DiscreteSpace([0.0, 0.25, 0.5, 0.75, 1.0])} @@ -328,7 +328,7 @@ def test_n_parallel_discrete(self): class TestBayesianSearchOptimize: def test_optimize_returns_correct_shapes(self): space = {"x": sp.DiscreteSpace([0.0, 0.25, 0.5, 0.75, 1.0])} - state, optimizer = bayesian.BayesianSearch.init(space, n_max=5, n_parallel=1) + state, optimizer = bayesian.BayesianSearch.init(space, n_max=5) func = lambda key, config: -(config["x"] ** 2) state, (params_hist, results_hist) = optimizer.optimize( state, jax.random.PRNGKey(0), func @@ -338,14 +338,14 @@ def test_optimize_returns_correct_shapes(self): def test_optimize_fills_state(self): space = {"x": sp.DiscreteSpace([0.0, 0.25, 0.5, 0.75, 1.0])} - state, optimizer = bayesian.BayesianSearch.init(space, n_max=5, n_parallel=1) + state, optimizer = bayesian.BayesianSearch.init(space, n_max=5) func = lambda key, config: -(config["x"] ** 2) state, _ = optimizer.optimize(state, jax.random.PRNGKey(0), func) assert int(state.mask.sum()) == 5 def test_optimize_finds_optimum(self): space = {"x": sp.DiscreteSpace([0.0, 0.25, 0.5, 0.75, 1.0])} - state, optimizer = bayesian.BayesianSearch.init(space, n_max=20, n_parallel=1) + state, optimizer = bayesian.BayesianSearch.init(space, n_max=20) func = lambda key, config: -(config["x"] ** 2) state, (params_hist, results_hist) = optimizer.optimize( state, jax.random.PRNGKey(0), func @@ -355,7 +355,7 @@ def test_optimize_finds_optimum(self): def test_optimize_with_array_result(self): space = {"x": sp.DiscreteSpace([0.0, 0.25, 0.5, 0.75, 1.0])} - state, optimizer = bayesian.BayesianSearch.init(space, n_max=3, n_parallel=1) + state, optimizer = bayesian.BayesianSearch.init(space, n_max=3) func = lambda key, config: jnp.array([-(config["x"] ** 2)]) state, (_, results_hist) = optimizer.optimize( state, jax.random.PRNGKey(0), func @@ -365,7 +365,7 @@ def test_optimize_with_array_result(self): def test_optimize_converges_toward_optimum(self): space = {"x": sp.DiscreteSpace([0.0, 0.25, 0.5, 0.75, 1.0])} state, optimizer = bayesian.BayesianSearch.init( - space, n_max=20, n_parallel=1, acquisition=acq.UCB() + space, n_max=20, acquisition=acq.UCB() ) func = lambda key, config: -(config["x"] ** 2) state, _ = optimizer.optimize(state, jax.random.PRNGKey(0), func) @@ -373,7 +373,7 @@ def test_optimize_converges_toward_optimum(self): def test_optimize_continuous_space(self): space = {"x": sp.LinearSpace(0.0, 1.0)} - state, optimizer = bayesian.BayesianSearch.init(space, n_max=5, n_parallel=1) + state, optimizer = bayesian.BayesianSearch.init(space, n_max=5) func = lambda key, config: -(config["x"] ** 2) state, (params_hist, results_hist) = optimizer.optimize( state, jax.random.PRNGKey(0), func @@ -384,7 +384,7 @@ def test_optimize_continuous_space(self): def test_optimize_continuous_with_ei_uses_observed_y_max(self): space = {"x": sp.LinearSpace(0.0, 1.0), "y": sp.LinearSpace(0.0, 1.0)} state, optimizer = bayesian.BayesianSearch.init( - space, n_max=20, n_parallel=1, acquisition=acq.EI() + space, n_max=20, acquisition=acq.EI() ) state = optimizer.update_state( state, jax.random.PRNGKey(0), jnp.array([100.0]), jnp.array([[0.5, 0.5]]) @@ -394,10 +394,10 @@ def test_optimize_continuous_with_ei_uses_observed_y_max(self): def test_optimize_minimize(self): space = {"x": sp.DiscreteSpace([0.0, 0.25, 0.5, 0.75, 1.0])} - state, optimizer = bayesian.BayesianSearch.init(space, n_max=20, n_parallel=1, maximize=False) + state, optimizer = bayesian.BayesianSearch.init(space, n_max=10, maximize=False) func = lambda key, config: config["x"] ** 2 state, _ = optimizer.optimize(state, jax.random.PRNGKey(0), func) - assert int(state.mask.sum()) == 20 + assert int(state.mask.sum()) == 10 assert float(optimizer.best_result(state)) == pytest.approx(0.0) def test_optimize_n_parallel_fills_buffer(self): @@ -458,7 +458,7 @@ def test_best_params_minimize(self): assert float(params["x"]) == pytest.approx(0.5) def test_best_result_after_full_optimize(self): - state, optimizer = bayesian.BayesianSearch.init(self.space, n_max=10, n_parallel=1) + state, optimizer = bayesian.BayesianSearch.init(self.space, n_max=10) func = lambda key, config: -(config["x"] ** 2) state, _ = optimizer.optimize(state, jax.random.PRNGKey(0), func) assert float(optimizer.best_result(state)) == pytest.approx( @@ -509,7 +509,7 @@ def test_log_length_scale_unchanged_with_single_observation(self): def test_tuned_length_scale_used_in_gp(self): state, optimizer = bayesian.BayesianSearch.init( - self.space, n_max=10, n_parallel=1, n_hparam_steps=20 + self.space, n_max=10, n_hparam_steps=20 ) func = lambda key, config: -(config["x"] ** 2) state, _ = optimizer.optimize(state, jax.random.PRNGKey(0), func) @@ -551,7 +551,7 @@ def test_get_next_params_mixed(self): "lr": sp.LogSpace(1e-4, 1e-1), "layers": sp.DiscreteSpace([1, 2, 3, 4]), } - state, optimizer = bayesian.BayesianSearch.init(space, n_max=20, n_parallel=1) + state, optimizer = bayesian.BayesianSearch.init(space, n_max=20) params = optimizer.get_next_params(state, jax.random.PRNGKey(0)) assert "lr" in params and "layers" in params assert params["lr"].shape == (1,) @@ -593,7 +593,7 @@ def test_two_arg_passes(self): class TestBayesianSearchOptimizeScan: def test_optimize_scan_runs(self): space = {"x": sp.LinearSpace(0.0, 1.0)} - state, optimizer = bayesian.BayesianSearch.init(space, n_max=5, n_parallel=1) + state, optimizer = bayesian.BayesianSearch.init(space, n_max=5) func = lambda key, config: -(config["x"] ** 2) state, (params_hist, results_hist) = optimizer.optimize_scan( state, jax.random.PRNGKey(0), func @@ -719,9 +719,9 @@ def test_lbfgs_improves_over_seed(self): state, key, jnp.array([r]), jnp.array([[x, y_]]) ) params = optimizer.get_next_params(state, key) - assert params["x"].shape == (optimizer.n_parallel,) - assert all(0.0 <= float(v) <= 1.0 for v in params["x"]) - assert all(0.0 <= float(v) <= 1.0 for v in params["y"]) + assert params["x"].shape == (1,) + assert 0.0 <= float(params["x"][0]) <= 1.0 + assert 0.0 <= float(params["y"][0]) <= 1.0 class TestKrigingBelieverHallucination: @@ -749,10 +749,10 @@ def _make_state_with_obs( state = optimizer.update_state(state, key, jnp.array([float(i)]), x) return state, optimizer - def test_default_hallucination_is_sample(self): + def test_default_hallucination_is_mean(self): space = {"x": sp.LinearSpace(0.0, 1.0)} _, optimizer = bayesian.BayesianSearch.init(space) - assert isinstance(optimizer.hallucination, acq.SampleHallucination) + assert isinstance(optimizer.hallucination, acq.MeanHallucination) def test_mean_hallucination_optimize_runs(self): space = {"x": sp.LinearSpace(0.0, 1.0)}