[Installation] Failed to import spectral module

Hi,

Thanks for this amazing library!
I installed the module using the command given in the README ``pip install jax-cfd[complete]`` to have access to all submodules.

When I try to import the spectral module via `import jax_cfd.spectral as spectral` I get the following error message:

```python
import jax_cfd.spectral as spectral
Traceback (most recent call last):
File "", line 1, in
ModuleNotFoundError: No module named 'jax_cfd.spectral'
```

Some more information about my system and the environment:

- Python version: ``Python 3.9.7 (default, Sep 16 2021, 13:09:58) ``
- JAX version: ``0.2.21``
- JAX CFD version: ``0.1.0``
- OS: ``Linux 5.10.16.3-microsoft-standard-WSL2``

Is there any way that I can resolve this issue?

Thanks in advance!