Doing this enables general use of `jax` >= 0.2.0; however, we still need to find
a better fix for the shape-related problems introduced by omnistaging.

Closes #43.