Slow import with tensorflow_probability

[patel-zeel]

Description of the bug

Hi @wesselb, when I try to import stheno with tensorflow_probability (reproduced in a fresh Google colab environment), it takes a long time, around 30 seconds.

Code	Time taken
from time import time init = time() import tensorflow_probability.substrates.jax as tfp from stheno import GP, EQ, PseudoObs, PseudoObservationsFITC, PseudoObsDTC print(time()-init)	28 seconds
from time import time init = time() from stheno import GP, EQ, PseudoObs, PseudoObservationsFITC, PseudoObsDTC print(time()-init)	3.6 seconds
from time import time init = time() import tensorflow_probability.substrates.jax as tfp print(time()-init)	3.79 seconds