diff --git a/requirements.txt b/requirements.txt
index ba2cd2c..00827fc 100644
--- a/requirements.txt
+++ b/requirements.txt
@@ -6,7 +6,7 @@
 #
 absl-py==0.11.0           # via dm-haiku, jax, jaxlib
 dm-haiku==0.0.2           # via pandas_toolkit (setup.py)
-jax-toolkit==0.1.2        # via pandas_toolkit (setup.py)
+jax-toolkit==0.2.0        # via pandas_toolkit (setup.py)
 jax==0.2.5               # via jax-toolkit, pandas_toolkit (setup.py)
 jaxlib==0.1.56            # via jax-toolkit, pandas_toolkit (setup.py)
 joblib==0.17.0            # via scikit-learn