BUG: `pymc.__init__` should import `math`?

[thomasaarholt]

Describe the issue:

The type checker pyright doesn't find the pymc.math module when pymc is imported as pm. While it works at runtime, this is still a bug in pymc.__init__.py, according to the author of pyright.

Looking at pymc.__init__.py, I believe the intention is for pm. to expose most of the modules that are wildcard-imported from, so we should modify

from pymc import _version, gp, ode, sampling
from pymc.math import (...)

to

from pymc import _version, gp, ode, sampling, math # and any others
from pymc.math import (...)

Reproduceable code example:

# in vscode, with type checking on
import pymc as pm
foo = pm.math.dot
         ^^^^ red squiggly lines

Error message:

No response

PyMC version information:

installed 5.10.4 in a fresh python 3.11 venv on mac

Context for the issue:

No response

[ricardoV94]

Sounds good. I think we actually shouldn't be importing anything from math at the root level, but that's a separate issue