diff --git a/jax_cfd/__init__.py b/jax_cfd/__init__.py
index 9a28310..6569478 100644
--- a/jax_cfd/__init__.py
+++ b/jax_cfd/__init__.py
@@ -14,6 +14,6 @@
 
 """Defines the JAX-CFD module for computational fluid dynamics."""
 
-__version__ = '0.2.0'
+__version__ = '0.2.1'
 
 import jax_cfd.base
diff --git a/setup.py b/setup.py
index 60f5034..fe3d87f 100644
--- a/setup.py
+++ b/setup.py
@@ -22,7 +22,7 @@
 
 setuptools.setup(
     name='jax-cfd',
-    version='0.2.0',
+    version='0.2.1',
     license='Apache 2.0',
     author='Google LLC',
     author_email='noreply@google.com',