Port the Prior class from pymc-marketing into the pymc repo
#7416
Replies: 9 comments 33 replies
-
|
We already have a way to define variables outside of a model, what if we try to use that? x = pm.Normal.dist(0, 1)
with pm.Model() as m:
m.register_rv(x, name="x")Pros: allows as much expressivity as PyMC models and uses almost same API. Cons:
Alternative what about lambda functions / callables / functools.partial? Also may be worth checking the approach @jessegrabowski took with Statspace module, which is kind of "user defines a model that fills in the blanks, with defaults for things that are not specified". |
Beta Was this translation helpful? Give feedback.
-
|
I like the |
Beta Was this translation helpful? Give feedback.
-
|
The Prior class's main benefit is the explicit dims outside of a model context which then supports the broadcasting. PyMC requires a model context in order to specify dims to my knowledge. Also think additional benefits of class methods for various applications which won't come with partial. Those are:
Though we do have check on the strings with the available PyMC distributions, maybe use of distributions at initialization would be more explicit. i.e |
Beta Was this translation helpful? Give feedback.
-
|
Yeah those go beyond partial. Thanks. But any reason you can't do those with the The only one that's trickier is dims. We could allow dims outside since pymc supports symbolic shapes. We would just set them to the correct value when registering them in the model |
Beta Was this translation helpful? Give feedback.
-
|
If we add named dims to pytensor it would solve all of our problems :) |
Beta Was this translation helpful? Give feedback.
-
|
Yes, seems like dims in pytensor could power this functionality as well. Is there any scoping of that? What would be a realistic timeline to implement that? |
Beta Was this translation helpful? Give feedback.
-
|
long |
Beta Was this translation helpful? Give feedback.
-
|
Just a point on this:
An additional way is not necessarily a plus. Users have one more thing to learn / try to understand when they come across examples. This assumes an overlap between pymc users and those other libraries It's also limited in different ways that may lead us to more complexity or reinvent pymc v2 with strings: e.g
Re: the centered/non-centered is a specific thing that may make more sense as a |
Beta Was this translation helpful? Give feedback.
-
|
Also curious what @tomicapretto is doing in over bambi, and what are the pain points |
Beta Was this translation helpful? Give feedback.
-
|
We have had an abstraction for priors in Bambi since a long time ago, but it's much simpler than what Will implemented in pymc-marketing. In our case, it's just a structure that stores some data that we use to build PyMC distributions when we are writing the PyMC model. From what I can see, the Prior class in pymc-marketing does a lot more. I'll try to summarise my thoughts here. What I like
What I don't like
In my ideal world, users should be able to do this with PyMC or PyTensor variables outside a PyMC model. mu = pm.Normal()
sigma = pm.InvGamma()
beta = pm.Normal(mu, sigma, dims="geo")
beta_star = pm.NonCentered(beta)
beta_star.to_graphviz()With all that said, I do think the Prior class is a great addition to PyMC Marketing, that it makes the devs life easier and it can help users to build a PyMC model using components from PyMC Marketing. I'm just hesitant to include |
Beta Was this translation helpful? Give feedback.
-
|
Yeah, I agree with these cons. |
Beta Was this translation helpful? Give feedback.
-
|
Is this something like prior? class Prior:
def __init__(self, dist, name, *args):
... # save
@staticmethod
def _eval(arg):
if isinstance(arg, Prior):
return arg.dist(
arg.name,
*(_eval(arg) for arg in args),
**dict(kwarg=_eval(arg) for kwarg, arg in self.kwargs)
)
else:
return arg
def eval(self, coords=None, model=UNSET):
# Unset makes nested models
# None makes standalone models
with pm.Model(coords=coords, model=model) as m:
return _eval(self)
def prior(self, draws=1000):
return pm.draw(self.eval(model=None))
def plot(self, draws=1000):
return pm.plots.plot_dist(
self.prior(draws=draws)
) |
Beta Was this translation helpful? Give feedback.
-
Thats good. Is it exactly numpy standard? Are there parsing utilities? Does those exist in either numpy or pytensor |
Beta Was this translation helpful? Give feedback.
-
|
This is using numpy vectorize signature, but it's a PyTensor thing. You should have it as a property of |
Beta Was this translation helpful? Give feedback.
-
|
Nice. I was just exploring it a bit and created an issue here: pymc-labs/pymc-marketing#841 Explicit dims work but could probably be relaxed Was noticing this behavior of Categorical. Is this weird / a bug? p = pm.Dirichlet.dist(a=[1, 1, 1])
x = pm.Categorical.dist(p=p, shape=(5, 3)) # eval works
y = pm.Categorical.dist(p=p[:, None], shape=(5, 3)) # eval also works but unintuitive
z = pm.Categorical.dist(p=p[:, None], shape=(3, 5)) # eval doesn't work because of TypeError shape mismatchWould think padding p would make variable of shape=(3, 5) rather than (5, 3) but that causes shape mismatch |
Beta Was this translation helpful? Give feedback.
-
|
It behaves as expected even if it may be a bit surprising. Notice the shape of the last dimension of p = pm.Dirichlet.dist(a=np.ones(100))
x = pm.Categorical.dist(p=p, shape=(5, 3)) # (5, 3) categorical draws with values between 0-99Your x is totally fine, it's a (5,3) categorical draws with values between 0-2 p = pm.Dirichlet.dist(a=[1], shape=(3, 1))
y = pm.Categorical.dist(p=p, shape=(5, 3))
assert y.eval().max() == 0Your case is technically wrong at runtime, because the probabilities are not always 1 (which in this case also means, they don't add up to 1). In this case Categorical could raise an Error, but we don't. The behavior is undefined, and in this case can return (0, 1), even though 1 should not be possible. Your z variable is wrong, because of numpy broadcsating. You have an input with shape (3,1), the 1 (the last dimension) is a core dimension, so it doesn't matter for broadcasting. So you have a batch of (3,) p and are trying to use that for a variable with shape (3, 5), which doesn't match. It's similar to if you tried to run the following multinomial: pm.Multinomial.dist(n=np.ones((3,), p=np.ones(100)/100, shape=(3, 5,))Or pm.Multinomial.dist(n=np.ones((3,), p=np.full(shape=(3, 5), fill_value=1/5))
|
Beta Was this translation helpful? Give feedback.
-
|
Ah. Yeah, makes sense. See my mistake. |
Beta Was this translation helpful? Give feedback.
-
|
To add very little to this covo, |
Beta Was this translation helpful? Give feedback.
-
|
Do you have some links for where this takes place for reference @jessegrabowski |
Beta Was this translation helpful? Give feedback.
-
|
I had the idea that models should implement a default_priors property that would give a dictionary of var_name:prior pairs, which could be used by a (not yet implemented) method |
Beta Was this translation helpful? Give feedback.
-
|
@jessegrabowski your approach sounds a lot like pattern substitution. Indulge my use of S-expressions to represent nested callables transparently Model = (pm.Normal, "y", "mean", (pm.HalfNormal, "sigma", "std")) Then user needs to provide "mean" and "std" in this example or you could give defaults. model_config = { To build the model you start evaling the S-expressions, replacing place_holders by expressions in model_config. Prior does something like S expression (and can be used as an alternative), by allowing you to define a lazy distribution to plug-in in the model. I guess in STS you're doing it with clone_replace on some root/dummy variables? |
Beta Was this translation helpful? Give feedback.
-
|
Another thing worth discussing. Do we want this to allow defining intermediate graphs. Say library offers model with components A -> B -> C and implements A and C and user can customize B. I don't think any of the existing suggestions would work for this use case, but can perhaps be tweaked for it. Stuff like prior and plotting wouldn't be possible for B objects obviously |
Beta Was this translation helpful? Give feedback.
-
|
I don't really follow this. Do you have some pseudo-code or example in mind? |
Beta Was this translation helpful? Give feedback.
-
|
In a model like A = pm.Normal("A")
B = pm.Gamma("B", mu=A, dims="country")
C = pm.Norma("C", mu=B[idx], dims="trial", observed=data)Would it be useful for the user to control what B looks like? (I'm not sure) |
Beta Was this translation helpful? Give feedback.
-
|
Btw I think we definitely want to figure out something like this. Since there are 3 distinct libraries that would depend on this and there's probably some room for tweaking I would suggest adding it first to experimental and once all 3 libraries are happy with it and haven't felt the need to tweak it further we move to pymc? |
Beta Was this translation helpful? Give feedback.
-
|
Think this is fair solution. Definitely would like to build out the solution a bit more as it came pretty naturally from the dictionary configurations of pymc-marketing which didn't have much functionality to begin with. We just need to bring in pymc-experimental with pymc-marketing again. Any opposition to this @juanitorduz ? The dependencies seems to just be pymc and sklearn so there shouldn't be any conflicts or restrictions |
Beta Was this translation helpful? Give feedback.
-
|
I had something like this in the past. What I came up with is creating a closure that is provided by user and executed inside the model context from functools import partial
def priors_fn(**config) -> dict[str, pt.TensorVariable]:
...
class PriorsCls(pydantic.BaseModel):
...
def __call__(self) -> dict[str, pt.TensorVariable]:
...
def lazy_get(cfg: dict[K_t, V_t], key: str, on_empty: Callable[[], V_t], on_value: Callable[[V_t], V_t]=lambda x: x) -> V_t:
value = cfg.get(key)
if key is None:
value = on_empty()
else:
value = on_value(value)
return value
def lazy_get_deterministic(cfg: dict[str, pt.TensorValue], key: str, on_empty: Callable[[], pt.TensorValue], possible_dims: list[str]) -> pt.TensorValue:
return lazy_get(cfg, key, on_empty, lambda v: pm.Deterministic(key, v, dims=infer_dims(v, possible_dims=possible_dims)))
def model_builder(priors_fn):
with pm.Model(coords=...) as model:
with pm.Model("user_priors"):
priors = priors_fn()
cfg = partial(lazy_get, priors)
cfg_deterministic = partial(lazy_get_deterministic, priors)
b = cfg("b",
lambda: pm.Normal("b", dims="c"),
lambda v: pm.Deterministic("b", v, dims=infer_dims(v, possible_dims=["a", "b", "c"])) # infer_dims might be needed if dim number might be overriden
)
c = cfg_deterministic("g", lambda: pm.Normal("g", dims="c"), possible_dims=["a", "b", "c"])Effectively, the user facing protocol is the parametrized lambda function, with the only assumption to be called under a model context.
I believe some nice api over cfg("b",
lambda: pm.Normal("b", dims="c"),
lambda v: pm.Deterministic("b", v, dims=infer_dims(v, possible_dims=["a", "b", "c"]))
)can be also proposed, but the nature still seems single use inside the model and that is not a requirement to make use of it. What about def infer_dims(value, possible_dims: list[str]):
if value.ndim:
# slice from the end
return possible_dims[-value.ndim:]
else:
# we have scalar
return None |
Beta Was this translation helpful? Give feedback.
-
|
So it seems like the ability of having distributions outside of a model would be pretty useful then. The |
Beta Was this translation helpful? Give feedback.
-
I propose that we port the excellent
Priorclass by @wd60622 from thepymc-marketingrepo.What is it & what does it do?
The
Priorclass is basically a wrapper around pymc distributions. It has a number of useful properties:pymc, it is a way of allowing users to define custom priors in pre-build models, or model factory functions/classes. This supersedes a previous way of doing it based on parsing dicts. This is exactly what it's doing inpymc-marketing. I would definitely bring it intoCausalPy. And in theory it could also be used in packages likebambi- though I have no idea if there's a desire to do that.Bringing it in to
pymcwould allow an additional way to specify priors. This would be most appealing in hierarchical modeling situations, for example.More info
Priorapi docs in pymc-marketing: https://www.pymc-marketing.io/en/latest/api/generated/pymc_marketing.prior.html👆Please do check out these pages to get properly up to speed
Options
I see at least 4 possible options, and it would be great to get some input from folks.
pymc-marketing. This would then remain out of reach of other packages built on top of pymc.pymc-experimentalwith potential to move it intopymcin the futurepymcnow.(Personally, I like options 4, 2, and 3, in that order)
Beta Was this translation helpful? Give feedback.
All reactions