diff --git a/neural_tangents/__init__.py b/neural_tangents/__init__.py
index 4a841e4a..7d8a20e0 100644
--- a/neural_tangents/__init__.py
+++ b/neural_tangents/__init__.py
@@ -16,7 +16,7 @@
 """Public Neural Tangents modules and functions."""
 
 
-__version__ = '0.3.5'
+__version__ = '0.3.6'
 
 
 from neural_tangents import predict