search

wesselb / stheno

Gaussian process modelling in Python
MIT License
217 stars 18 forks source link

Slow import with tensorflow_probability #25

Closed patel-zeel closed 1 year ago

patel-zeel commented 2 years ago

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
wesselb commented 2 years ago

Hey @patel-zeel! Thanks for opening an issue.

I can confirm that the load times are indeed slow. My suspicion is that Plum is to blame here, which recently started using inspect to check whether or not from __future__ import annotations is loaded or not. This mechanism should be improved soon, at which point this performance issue should automatically disappear. Hence, for now, unless this is blocking you, I'll leave this open, and revise the issue once the mechanism is Plum has been improved.

wesselb commented 1 year ago

Hey @patel-zeel! With the new release of Plum, this issue should now be fixed.