Remove flax dependency

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", "jaxtyping"]
 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"
 ]

src/hyperoptax/base.py

@@ -1,10 +1,12 @@
+import dataclasses
 import inspect
 import warnings
+from dataclasses import dataclass
 from typing import Callable
 
 import jax
 import jax.numpy as jnp
-from flax import struct
+from jaxtyping import PyTree
 
 
 def _validate_func(func):
@@ -35,11 +37,15 @@ 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: PyTree
+
+    def replace(self, **kwargs) -> "OptimizerState":
+        return dataclasses.replace(self, **kwargs)
 
 
 class Optimizer:
@@ -55,7 +61,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[PyTree, jax.Array]]:
         """
         High Level API for optimizing a function over a space.
         Not recommended if you want to do fancy things
@@ -86,7 +92,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[PyTree, jax.Array]]:
         """
         Like optimize, but uses jax.lax.scan for the inner loop.
 
@@ -140,7 +146,7 @@ def update_state(
         state: OptimizerState,
         key: jax.random.PRNGKey,
         results: jax.Array,
-        params: struct.PyTreeNode = None,
+        params: PyTree | None = None,
     ) -> OptimizerState:
         """
         Updates the optimizer state based on the results of the function.
@@ -151,9 +157,9 @@ def get_next_params(
         self,
         state: OptimizerState,
         key: jax.random.PRNGKey,
-        params: struct.PyTreeNode = None,
+        params: PyTree | None = None,
         results: jax.Array = None,
-    ) -> struct.PyTreeNode:
+    ) -> PyTree:
         """
         Gets the next parameters to sample from the space.
         Returns a batched pytree where every leaf has shape (n_parallel, ...).

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`.
 

src/hyperoptax/grid.py

@@ -2,13 +2,14 @@
 
 import jax
 import jax.numpy as jnp
-from flax import struct
+from jaxtyping import PyTree
 
 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:
+    ) -> 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):

src/hyperoptax/random.py

@@ -2,7 +2,7 @@
 
 import jax
 import jax.numpy as jnp
-from flax import struct
+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,
-    ) -> struct.PyTreeNode:
+    ) -> PyTree:
         """Sample ``n_parallel`` independent configurations from the search space."""
         def sample_once(k):
             subkeys = utils.make_key_tree(state.space, k)

src/hyperoptax/spaces.py

@@ -1,23 +1,29 @@
+from abc import ABC, abstractmethod
+from dataclasses import dataclass, field
+
 import jax
 import jax.numpy as jnp
-from flax import struct
 
 
 # transformation between logs
 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
 
     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,18 +134,23 @@ 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
     integer hyperparameters whose scale spans orders of magnitude
     (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)

src/hyperoptax/utils.py

@@ -1,13 +1,13 @@
 import jax
-from flax import struct
+from jaxtyping import PyTree
 
 from hyperoptax import spaces as sp
 
 
 def make_key_tree(
-    pytree: struct.PyTreeNode,
+    pytree: PyTree,
     subkey: jax.random.PRNGKey,
-) -> struct.PyTreeNode:
+) -> PyTree:
     """Split ``subkey`` into a pytree of PRNGKeys matching the structure of ``pytree``.
 
     :class:`~hyperoptax.spaces.Space` objects are treated as leaves, so the

tests/test_grid.py

@@ -1,3 +1,4 @@
+
 import jax
 import jax.numpy as jnp
 import pytest

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():