bump version

Author: mariogeiger

ChangeLog.md

@@ -5,6 +5,8 @@ The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/),
 and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0.html).
 
 ## [Unreleased]
+
+## [0.20.1] - 2023-09-13
 ### Changed
 - Add back the optimizations with the lazy `._chunks` that was removed in 0.19.0
 

e3nn_jax/__init__.py

@@ -1,4 +1,4 @@
-__version__ = "0.20.0"
+__version__ = "0.20.1"
 
 from e3nn_jax._src.config import config
 from e3nn_jax._src.rotation import (

e3nn_jax/_src/utils/vmap.py

@@ -29,6 +29,7 @@ def vmap(
         >>> y = vmap(e3nn.scalar_activation)(x)
         >>> y.zero_flags
         (False, True)
+        >>> assert y.chunks[1] is None
     """
     if in_axes == -1:
         raise ValueError("in_axes=-1 is not supported by e3nn.vmap")