Problem of resolving tensorflow_probability distributions in Vscode


I am using

from tensorflow_probability import distributions as tfd

in vscode,

but the vscode resolves distributions as a variable with ModuleType instead of a Module,
so that it can not find subclasses of distributions.

By finding the definition of this distributions,
it points to tensorflow_probability/python/ line 82,
where distributions is given the types.ModuleType

It seems that those submodules of tensorflow_probability are lazily loaded,
but I am sorry I am not familiar with lazy loading.

I wonder how to import the distributions and other subclass of tensorflow_probability as normal module so that the IDE can find subclass of distributions for autocompletion.

Appreciate any helps!

The version of tensorflow is 2.10,
the version of tensorflow_probability is 0.18.0