Material Missing at import: strange...

Hi,
I am on Google Colab and
```python
! pip install git+https://github.com/deepmind/distrax.git
from distrax import ScalarAffine, Tanh, DiagLinear, Shift
```
I get this error
```
ImportError: cannot import name 'DiagLinear' from 'distrax' (/usr/local/lib/python3.7/dist-packages/distrax/__init__.py)
```

But also strange
```python
from distrax._src.bijectors.diag_linear import DiagLinear
from distrax._src.bijectors.linear import Linear
```
The first import is ok, but the second complains
```
ModuleNotFoundError: No module named 'distrax._src.bijectors.linear'
```

something wrong in my install??? I check that the version is (0.1.2)

Thanks