Open ricardoV94 opened 4 months ago
codecov[bot]
commented
4 months ago Attention: Patch coverage is 96.72131% with 2 lines in your changes are missing coverage. Please review.
Project coverage is 92.26%. Comparing base (
244fb97) to head (af027ad).
ricardoV94
commented
4 months ago Is there any expected advantage to this specialization besides avoiding constraints checks (which the user can already do via the model check_bounds flag)?
aseyboldt
commented
4 months ago Nice :D
In the case of the ZeroSumNormal probably not really much, once we figure out the expression in the untransformed space it doesn't change too much, and I guess it is nice to have for instance so that the logp can be computed. Just being able to do it on the transformed space would have saved us doing some math...
I think it is quite common of the logp expressions to be a bit cleaner and possibly numerically more stable on the transformed space.
For instance for the beta distribution we have this here in the code currently (before the bounds check):
res = (
pt.switch(pt.eq(alpha, 1.0), 0.0, (alpha - 1.0) * pt.log(value))
+ pt.switch(pt.eq(beta, 1.0), 0.0, (beta - 1.0) * pt.log1p(-value))
- (pt.gammaln(alpha) + pt.gammaln(beta) - pt.gammaln(alpha + beta))
)
res = pt.switch(pt.bitwise_and(pt.ge(value, 0.0), pt.le(value, 1.0)), res, -np.inf)on the (logit) transformed space it could simply be
normalizing_factor = pt.gammaln(alpha + beta) - pt.gammaln(alpha) - pt.gammaln(beta)
density = - pt.log1p(pt.exp(-logit_x)) * (alpha + beta) - logit_x * beta + normalizing_factorAll the switches are gone, the graph is much shorter (in the original value is actually expit(x), so we actually have things like log(expit(x)) and log(1- expit(x)) in the graph).
Or for instance for the gamma with log transform: Now, we have
beta = pt.reciprocal(scale)
res = -pt.gammaln(alpha) + logpow(beta, alpha) - beta * value + logpow(value, alpha - 1)
res = pt.switch(pt.ge(value, 0.0), res, -np.inf)Which could be
beta = pt.reciprocal(scale)
norm = - beta * pt.log(alpha) - pt.gammaln(beta)
density = beta * log_x - exp(log_x) / alpha(I guess we could also pass in the transformed point into the logp function, so that the exp(log_x) could just be x. The original had things like logpow(exp(x), alpha - 1) in the graph, the new one doesn't. Not a big deal I think, but it is cleaner.
For the logit-normal what we have to do right now is even a bit silly:
res = pt.switch(
pt.or_(pt.le(value, 0), pt.ge(value, 1)),
-np.inf,
(
-0.5 * tau * (logit(value) - mu) ** 2
+ 0.5 * pt.log(tau / (2.0 * np.pi))
- pt.log(value * (1 - value))
),
)This contains logit(value). So we first compute value = expit(logit_x) and then just undo it. We also have the - pt.log(value * (1 - value)), which is just the jacobian det of the transform. So we are taking a point in the untransformed space, compute the transformed value and the jac det, then we (manually) untransform the value again, and (manually) subtract the logdet again...
I guess things like this might also be happining in the dirichlet dists, but looking through that is a bit more work... Nothing here really is a game-changer, but I think it might make quite a few graphs a bit cleaner and a bit more stable.
There is also an additional thing we could do with this that right now we can't do with transformations at all. I'm not 100% sure if we should do this though:
Right now all our transformations are injective, because the whole trick with the jacobian determinant doesn't work otherwise. But if we can compute the logp on the transformed space, we could get rid of that requirement. So for instance we could have a Horseshoe distribution with a transformed space that contains the lambda and the x values, and the tranformation just multiplies those.
Or we could have a transformation that maps a point in 2d space to an angel, and use that in for instance the VanMises distribution to avoid the topology problems we have there right now.
ricardoV94
commented
4 months ago Thanks for the reply
Re: switches/bound checks, most of those could probably be removed with some domain analysis like exp(x) -> positive so any ge(x, 0) are true, that could be useful beyond PyMC. This could be coupled with user provided type hints (which we are kind of doing with transformed values here) to seed initial info that is not extractable from the graph.
Re: examples, many of those sound like interesting rewrites that PyTensor already does or could do. That need no be a blocker to allow users to work directly on transformed space if they want, but may be a good source of ideas for future rewrites.
The best argument for the functionality in my perspective is the "avoiding hard math", which is more than a nice to have, and could speedup development or allow models that are otherwise impossible to write without a hacky Potential?
For our code-base I think we still want to try and provide constrained space logps as much as possible, since not all we do (or may want to do) is logp-based sampling.
twiecki
commented
4 months ago What does this do?
aseyboldt
commented
4 months ago I just had a look at the generated graphs of each of those. Let's say we want the logp of the model:
with pm.Model(check_bounds=False) as model:
pm.ZeroSumNormal("x", shape=(10, 10), n_zerosum_axes=2)The logp + grad with the PR:
>>> func1 = model.logp_dlogp_function(mode="NUMBA")
>>> pytensor.dprint(func1._pytensor_function)
Sum{axes=None} [id A] 1
ββ Composite{((-0.5 * sqr(i0)) - 0.9189385332046727)} [id B] 'x_zerosum___logprob' 0
ββ x_zerosum__ [id C]
Neg [id D] 'x_zerosum___grad' 2
ββ x_zerosum__ [id C]
Inner graphs:
Composite{((-0.5 * sqr(i0)) - 0.9189385332046727)} [id B]
β sub [id E] 'o0'
ββ mul [id F]
β ββ -0.5 [id G]
β ββ sqr [id H]
β ββ i0 [id I]
ββ 0.9189385332046727 [id J]And without:
I think I like this PR :D
Maybe we can also add an implementation for the beta and the logit-normal dists:
@pm.logprob.abstract._transformed_logprob.register(pt.random.basic.BetaRV, pm.distributions.transforms.LogOddsTransform)
def transformed_beta_logp(op, transform, unconstrained_value, rv_inputs):
*_, alpha, beta = rv_inputs
logit_x = unconstrained_value
normalizing_factor = pt.gammaln(alpha + beta) - pt.gammaln(alpha) - pt.gammaln(beta)
logp = normalizing_factor - pt.log1p(pt.exp(-logit_x)) * (alpha + beta) - logit_x * beta
return pm.distributions.multivariate.check_parameters(
logp,
alpha > 0,
beta > 0,
"Alpha and beta parameters must be positive in Beta distribution",
)
@pm.logprob.abstract._transformed_logprob.register(pm.distributions.continuous.LogitNormalRV, pm.distributions.transforms.LogOddsTransform)
def transformed_logit_normal_logp(op, transform, unconstrained_value, rv_inputs):
*_, mu, sigma = rv_inputs
return pm.logp(pm.Normal.dist(mu, sigma), unconstrained_value)@aseyboldt perhaps a more fair graph comparison would be with Model(check_bounds=False). The mean switch for instance is just a bound check thing that could equally be removed in a model with the transform and the standard logp
aseyboldt
commented
4 months ago That was with check_bounds=False. For most dists that only disables the checks on the parameters, not the domain of the value. But I don't think that's the main source for the difference either.
But this comparison wasn't entirely fair, as the transformation wasn't needed at all anymore with the PR, while in most real applications that would still be needed when the variable is used downstream in the model. I guess this one is a fairer comparision:
with pm.Model(check_bounds=False) as model:
x = pm.ZeroSumNormal("x", shape=(100, 100), n_zerosum_axes=2)
pm.Normal("y", mu=x, sigma=1, shape=(100, 100))With PR:
Without PR
That's 140 vs 200 lines. Got to admit that I'm not 100% sure what the difference really is though...
It's also a bit faster, with 70ΞΌs vs 90ΞΌs.
ricardoV94
commented
4 months ago That was with check_bounds=False
Then I'm surprised the isclose(mean, 0) switch was not removed
aseyboldt
commented
4 months ago You are right, it seems check_bounds=False didn't do anything, because I wasn't compiling the pytensor functions in a model context. I had no idea I had to do that...
Description
Adds a specialized dispatch for transformed logps
TODO
Related Issue
Checklist
Type of change
π Documentation preview π: https://pymc--7188.org.readthedocs.build/en/7188/