diff --git a/checkpoint/pyproject.toml b/checkpoint/pyproject.toml
index 68eb7dd2..96955295 100644
--- a/checkpoint/pyproject.toml
+++ b/checkpoint/pyproject.toml
@@ -50,5 +50,6 @@ testing = [
     'mock',
     'flax',
     'pytest',
-    'pytest-xdist'
+    'pytest-xdist',
+    'chex',
 ]