Type-hint domain and parameterization information for parameters

[williambdean]

I made a distribution explorer tool which is based on Annotated type hints which I think would be cool to bring to PyMC eco.

It doesn't have to be done the same way, but Annotated type hints were pretty easy to work and encode arbitrary metadata

from typing import Annotated

Real = Annotated[float, "Real"]
PositiveReal = Annotated[float, "PositiveReal"]


def foo(mu: Real, sigma: PositiveReal): 
    ...

mu_type = foo.__annotations__["mu"] # typing.Annotated[float, 'Real']
mu_type.__metadata__ # ('Real',)

In my solution, I didn't tackle different parameterization but I suspect it could be done with a similar approach.

def foo(mu: Real, sigma: Annotated[PositiveReal, 1], tau: Annotated[PositiveReal, 2]):
  ...

foo.__annotations__
"
{'mu': typing.Annotated[float, 'Real'],
 'sigma': typing.Annotated[float, 'PositiveReal', 1],
 'tau': typing.Annotated[float, 'PositiveReal', 2]}
"

Would love to hear thoughts on this idea!

[jessegrabowski]

I would like if this information actually lived on the compute graph, not just in the typehints. Symbolic reasoning about metadata has the potential to be extremely powerful and enable some cool rewrites.

In my mind this is closely related to discussion like pymc-devs/pytensor#573 and pymc-devs/pymc#7812, where we wish we had access to some metadata about the symbolic variables (structured matrices in the first, units of measurement in the second) to do some smart rewrites. Information about the domain of a variable would also fall into the category. For example, if we saw something like max(x, 0) and x is tagged strictly positive, we could rewrite away the max operation, or do log(exp(x)) = x and sqrt(sqr(x)) without a switch.

[aloctavodia]

That information could be useful in many ways, for instance, ArviZ-plots could use automatically choose a visualization or warn the user of misuse of some test like a PPC not valid for binary data, etc. In PreliZ there is information about the support and some of that info I used for example https://preliz.readthedocs.io/en/latest/examples/gallery/predictive_explorer.html, but as @jessegrabowski points out, its probably better to have that info as part of the graph, otherwise it is difficult to know what happens once we apply some transformation.

[williambdean]

Yes, not just type hints. In my mind, there is some information that should be attached to the Distributions and some information to the variables. And there might be an overlap or transfer of this information.

How do you envision someone adding this information with the fact that distributions work with floats and other non-Variable types?