diff --git a/pyproject.toml b/pyproject.toml
index 2818b9a..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", "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"
 ]
diff --git a/src/hyperoptax/base.py b/src/hyperoptax/base.py
index 91d3300..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 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, ...).
diff --git a/src/hyperoptax/bayesian.py b/src/hyperoptax/bayesian.py
index b192f7b..e9f20e5 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`.
 
diff --git a/src/hyperoptax/grid.py b/src/hyperoptax/grid.py
index 17e62f2..401e6b3 100644
--- a/src/hyperoptax/grid.py
+++ b/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):
diff --git a/src/hyperoptax/random.py b/src/hyperoptax/random.py
index 9345bc2..af4964e 100644
--- a/src/hyperoptax/random.py
+++ b/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)
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..efa6b30 100644
--- a/src/hyperoptax/utils.py
+++ b/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
diff --git a/tests/test_grid.py b/tests/test_grid.py
index a949df5..03a0729 100644
--- a/tests/test_grid.py
+++ b/tests/test_grid.py
@@ -1,3 +1,4 @@
+
 import jax
 import jax.numpy as jnp
 import pytest
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():