Bounded variable transformation when not using HMC?

Questions version agnostic
6

Thanks for the helpful discussion @ricardoV94. Really appreciate the help and insight on automatic transforms. (FWIW I’m mainly using slice sampler.)

I’ve come up against an issue when working with non-transformed models that do have bounded variables. It’s immediately following this thread so I thought I’d write here, but happy to take the conversation somewhere else if that’d be preferred.

I’ve noticed that if a variable has a bounded distribution and is not transformed, the sampler may sometimes send out-of-bounds parameter values to the likelihood function. In particular what seems to be going on is,

  1. Step method proposes, or tests, a parameter value that is outside the bounds of that variable (this is fine and expected sometimes)
  2. pymc/pytensor calculates that the logp of that particular factor is -inf
  3. But pymc still cranks out the full sum of model logp terms, which means that the variable with logp=-inf gets passed on to the model likelihood term. (or in more complex models, it could just be that the impossible value of the hyperprior gets passed on to some other node in the DAG)

This is not an issue when the default transforms are used, because any possible real number that the step method could propose gets transformed to that variable’s bounds. In my case, this basically means that I run up against the same exception from my likelihood function as we were discussing above.

With the built-in pymc distributions, it looks like the logp functions use switch to ensure that they are robust to any parameter value. But in my case, the black-box likelihood may fail if parameter values are too extreme. I had thought that specifying bounded prior distributions would be sufficient to ensure that out-of-bounds values do not get passed along the DAG when computing model logp, but currently this is not the case.

Here’s a MWE of what I mean:

import numpy as np
import pymc as pm
import pytensor

# Simple regression on one slope term. Let's enforce that the slope variable
# remains within the following bounds.
bounds = (-2, 2)

# Data generation
X = np.random.uniform(size=1000)
m = 1
assert bounds[0] <= m <= bounds[1]
Y = m * X + np.random.randn(1000) * 0.1

# Model using an extended Normal logp that raises AssertionError if its coefficient
# is outside of the enforced bounds
def regression_logp(x, y, beta, sigma):
    b = pytensor.raise_op.Assert("unexpected beta!")(
        beta,
        pytensor.tensor.all(
            [pytensor.tensor.ge(beta, bounds[0]), pytensor.tensor.le(beta, bounds[1])]
        ),
    )
    val = x * b
    return pm.Normal.logp(value=y, mu=val, sigma=sigma)

with pm.Model() as m:
    b = pm.TruncatedNormal("b", mu=0, sigma=5, lower=bounds[0], upper=bounds[1], transform=None)
    y_obs = pm.Potential("y_obs", regression_logp(x=X, y=Y, beta=b, sigma=0.1))

    # Sampling will fail due to AssertionError. This occurs with any step method.
    # But sampling doesn't fail if we keep the default transform on b.
    idata = pm.sample()

I’m wondering if it could be reasonable to modify the behavior of the pymc model’s logp() function. What if, rather than calculating all of the logp terms (in this line), we stop if any of the terms is -inf? I’m not sure if this would have any bad side-effects that I can’t think of right now. Or how it should be implemented—would we want a new pytensor accumulator Op that stops if a certain condition is met?