diff --git a/.github/workflows/tests_linters.yml b/.github/workflows/tests_linters.yml
index 3b6eee9..b78d8b1 100644
--- a/.github/workflows/tests_linters.yml
+++ b/.github/workflows/tests_linters.yml
@@ -4,10 +4,10 @@ on: [ push, pull_request ]
 
 jobs:
   tests-and-linters:
-    name: "Python 3.9 on GitHub Hosted runner"
+    name: "Python 3.10 on GitHub Hosted runner"
     runs-on: ubuntu-latest
     container:
-        image: python:3.9
+        image: python:3.10
 
     steps:
       - name: Install dependencies for viewer test
diff --git a/flashbax/buffers/sum_tree.py b/flashbax/buffers/sum_tree.py
index 9967d8e..24e64d9 100644
--- a/flashbax/buffers/sum_tree.py
+++ b/flashbax/buffers/sum_tree.py
@@ -12,28 +12,24 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from typing import TYPE_CHECKING, Optional, Tuple, Union
-
-if TYPE_CHECKING:  # https://github.com/python/mypy/issues/6239
-    from dataclasses import dataclass
-else:
-    from flax.struct import dataclass
+from typing import Optional, Tuple, Union
 
 import chex
 import jax
 import jax.numpy as jnp
 import numpy as np
-from flax import struct
 from jax import Array
 
+from flashbax.dataclass import dataclass, field
+
 
 @dataclass
 class SumTreeState:
     nodes: Array
     max_recorded_priority: Array
-    tree_depth: int = struct.field(pytree_node=False)
-    capacity: int = struct.field(pytree_node=False)
-    dtype: jnp.dtype = struct.field(pytree_node=False)
+    tree_depth: int = field(pytree_node=False)
+    capacity: int = field(pytree_node=False)
+    dtype: jnp.dtype = field(pytree_node=False)
 
 
 def get_tree_depth(capacity: int) -> int:
diff --git a/flashbax/buffers/sum_tree_test.py b/flashbax/buffers/sum_tree_test.py
index 80ba7d6..997fb87 100644
--- a/flashbax/buffers/sum_tree_test.py
+++ b/flashbax/buffers/sum_tree_test.py
@@ -255,8 +255,9 @@ def test_set_batch_scan_matches_set_batch_bincount(
     values1 = jax.random.permutation(rng_key1, values)
     values2 = jax.random.permutation(rng_key2, values)
 
-    assert jnp.all(indexes1 != indexes2)
-    assert jnp.all(values1 != values2)
+    # We only need the shuffles to differ somewhere to validate ordering invariance.
+    assert jnp.any(indexes1 != indexes2)
+    assert jnp.any(values1 != values2)
 
     state_bincount = init_state  # Re-init the state.
 
diff --git a/flashbax/dataclass.py b/flashbax/dataclass.py
new file mode 100644
index 0000000..cff227a
--- /dev/null
+++ b/flashbax/dataclass.py
@@ -0,0 +1,214 @@
+# Copyright 2023 InstaDeep Ltd. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+# Copyright 2024 The Flax Authors.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Utilities for defining custom classes that can be used with jax transformations."""
+
+import dataclasses
+import functools
+from collections.abc import Callable
+from typing import TypeVar, Union, overload
+
+import jax
+from typing_extensions import dataclass_transform  # pytype: disable=not-supported-yet
+
+_T = TypeVar("_T")
+
+
+def field(pytree_node=True, *, metadata=None, **kwargs):
+    return dataclasses.field(
+        metadata=(metadata or {}) | {"pytree_node": pytree_node}, **kwargs
+    )
+
+
+@dataclass_transform(field_specifiers=(field,))  # type: ignore[literal-required]
+@overload
+def dataclass(clz: _T, **kwargs) -> _T:
+    ...
+
+
+@dataclass_transform(field_specifiers=(field,))  # type: ignore[literal-required]
+@overload
+def dataclass(**kwargs) -> Callable[[_T], _T]:
+    ...
+
+
+@dataclass_transform(field_specifiers=(field,))  # type: ignore[literal-required]
+def dataclass(
+    clz: Union[_T, None] = None,
+    **kwargs,
+) -> Union[_T, Callable[[_T], _T]]:
+    """Create a class which can be passed to functional transformations.
+
+    .. note::
+      Inherit from ``PyTreeNode`` instead to avoid type checking issues when
+      using PyType.
+
+    Jax transformations such as ``jax.jit`` and ``jax.grad`` require objects that are
+    immutable and can be mapped over using the ``jax.tree_util`` methods.
+    The ``dataclass`` decorator makes it easy to define custom classes that can be
+    passed safely to Jax. Define JAX data as normal attribute fields, and use
+    ``pytree_node=False`` to define static metadata.
+
+    See example::
+
+      >>> from flax import struct
+      >>> import jax
+      >>> from typing import Any, Callable
+
+      >>> @struct.dataclass
+      ... class Model:
+      ...   params: Any
+      ...   # use pytree_node=False to indicate an attribute should not be touched
+      ...   # by Jax transformations.
+      ...   apply_fn: Callable = struct.field(pytree_node=False)
+
+      ...   def __apply__(self, *args):
+      ...     return self.apply_fn(*args)
+
+      >>> params = {}
+      >>> params_b = {}
+      >>> apply_fn = lambda v, x: x
+      >>> model = Model(params, apply_fn)
+
+      >>> # model.params = params_b  # Model is immutable. This will raise an error.
+      >>> model_b = model.replace(params=params_b)  # Use the replace method instead.
+
+      >>> # This class can now be used safely in Jax to compute gradients w.r.t. the
+      >>> # parameters.
+      >>> model = Model(params, apply_fn)
+      >>> loss_fn = lambda model: 3.
+      >>> model_grad = jax.grad(loss_fn)(model)
+
+    Note that dataclasses have an auto-generated ``__init__`` where
+    the arguments of the constructor and the attributes of the created
+    instance match 1:1. If you desire a "smart constructor", for example to
+    optionally derive some of the attributes from others,
+    make an additional static or class method. Consider the following example::
+
+      >>> @struct.dataclass
+      ... class DirectionAndScaleKernel:
+      ...   direction: jax.Array
+      ...   scale: jax.Array
+
+      ...   @classmethod
+      ...   def create(cls, kernel):
+      ...     scale = jax.numpy.linalg.norm(kernel, axis=0, keepdims=True)
+      ...     direction = direction / scale
+      ...     return cls(direction, scale)
+
+    Args:
+      clz: the class that will be transformed by the decorator.
+      **kwargs: arguments to pass to the dataclass constructor.
+
+    Returns:
+      The new class.
+    """
+    # Support passing arguments to the decorator (e.g. @dataclass(kw_only=True))
+    if clz is None:
+        return functools.partial(dataclass, **kwargs)  # type: ignore[bad-return-type]
+
+    # check if already a flax dataclass
+    if "_flax_dataclass" in clz.__dict__:
+        return clz
+
+    if "frozen" not in kwargs.keys():
+        kwargs["frozen"] = True
+    data_clz = dataclasses.dataclass(**kwargs)(clz)  # type: ignore
+    meta_fields = []
+    data_fields = []
+    for field_info in dataclasses.fields(data_clz):
+        is_pytree_node = field_info.metadata.get("pytree_node", True)
+        if is_pytree_node:
+            data_fields.append(field_info.name)
+        else:
+            meta_fields.append(field_info.name)
+
+    def replace(self, **updates):
+        """Returns a new object replacing the specified fields with new values."""
+        return dataclasses.replace(self, **updates)
+
+    data_clz.replace = replace
+
+    jax.tree_util.register_dataclass(data_clz, data_fields, meta_fields)
+
+    # add a _flax_dataclass flag to distinguish from regular dataclasses
+    data_clz._flax_dataclass = True  # type: ignore[attr-defined]
+
+    return data_clz  # type: ignore
+
+
+TNode = TypeVar("TNode", bound="PyTreeNode")
+
+
+@dataclass_transform(field_specifiers=(field,))  # type: ignore[literal-required]
+class PyTreeNode:
+    """Base class for dataclasses that should act like a JAX pytree node.
+
+    See ``flax.struct.dataclass`` for the ``jax.tree_util`` behavior.
+    This base class additionally avoids type checking errors when using PyType.
+
+    Example::
+
+      >>> from flax import struct
+      >>> import jax
+      >>> from typing import Any, Callable
+
+      >>> class Model(struct.PyTreeNode):
+      ...   params: Any
+      ...   # use pytree_node=False to indicate an attribute should not be touched
+      ...   # by Jax transformations.
+      ...   apply_fn: Callable = struct.field(pytree_node=False)
+
+      ...   def __apply__(self, *args):
+      ...     return self.apply_fn(*args)
+
+      >>> params = {}
+      >>> params_b = {}
+      >>> apply_fn = lambda v, x: x
+      >>> model = Model(params, apply_fn)
+
+      >>> # model.params = params_b  # Model is immutable. This will raise an error.
+      >>> model_b = model.replace(params=params_b)  # Use the replace method instead.
+
+      >>> # This class can now be used safely in Jax to compute gradients w.r.t. the
+      >>> # parameters.
+      >>> model = Model(params, apply_fn)
+      >>> loss_fn = lambda model: 3.
+      >>> model_grad = jax.grad(loss_fn)(model)
+    """
+
+    def __init_subclass__(cls, **kwargs):
+        dataclass(cls, **kwargs)  # pytype: disable=wrong-arg-types
+
+    def __init__(self, *args, **kwargs):
+        # stub for pytype
+        raise NotImplementedError
+
+    def replace(self: TNode, **overrides) -> TNode:
+        # stub for pytype
+        raise NotImplementedError
diff --git a/pyproject.toml b/pyproject.toml
index 71eaa3c..a0b1ae2 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -40,7 +40,6 @@ classifiers=[
     ]
 keywords=["reinforcement-learning", "python", "jax", "memory"]
 dependencies = [
-    'flax>=0.6.11',
     'chex>=0.1.8',
     'jax>=0.4.25',
     'jaxlib>=0.4.20',
@@ -66,7 +65,10 @@ dev = [
     'pre-commit>=2.20.0',
     'pytest>=7.4.2',
     'pytest-cov>=4.00',
-    'pytest-xdist>=3.0.2'
+    'pytest-xdist>=3.0.2',
+    'orbax-checkpoint==0.11.20',
+    'etils==1.13.0',
+    'importlib-resources==6.5.2',
     ]
 examples = [
     'distrax',