Material Missing at import: strange...

[jecampagne]

Hi,
I am on Google Colab and

! 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

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