diff --git a/parallax/__init__.py b/parallax/__init__.py
index a154b19..000e64c 100644
--- a/parallax/__init__.py
+++ b/parallax/__init__.py
@@ -14,6 +14,8 @@
 
 """Parallax library."""
 
+__version__ = "0.0.2"
+
 from flax import nnx
 import jax
 from parallax import offload
diff --git a/pyproject.toml b/pyproject.toml
index e073bc0..ddbd9ac 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -4,7 +4,7 @@ build-backend = "hatchling.build"
 
 [project]
 name = "google-parallax"
-version = "0.0.1"
+version = "0.0.2"
 authors = [
   { name="Jeff Carpenter", email="code@jeffcarp.com" },
   { name="Stella Yan", email="sueyan@gmail.com" },
@@ -20,11 +20,13 @@ license = {file = "LICENSE"}
 dependencies = [
   "flax",
   "jax",
+  "jaxtyping",
+  "numpy",
 ]
 
 [tool.hatch.build]
 include = [
-  "src/*"
+  "parallax/*"
 ]
 exclude = [
   "**/*_test.py"