Using CustomDist as a stochastic variable - causing TensorType error

Questions v5
1

Hello, I am a bit puzzled by the behaviour of CustomDist and am looking for some input.

Background: I am trying to model a set of data as being drawn from a well-behaved intrinsic distribution where the observations have added noise. I usually do this with a Gaussian, so I have a model

v ~ N(mu, sigma^2)
v_obs ~ N(v, sigma_obs^2)

with Gaussian priors on mu and sigma and known sigma_obs. I now wanted to extend this to a generalised Gaussian and since that is not a built-in distribution, I implemented a CustomDist with the following logp & random functions:

def logp(value: TensorVariable, mu: TensorVariable,
         alpha: TensorVariable, beta: TensorVariable) -> TensorVariable:
    
    lp = gennorm.pdf(value, beta, loc=mu, scale=alpha)
    return lp

# Function to draw a random value from the distribution
def random(
    mu: np.ndarray | float,
    alpha: np.ndarray | float,
    beta: np.ndarray | float,
    rng: Optional[np.random.Generator] = None,
    size : Optional[Tuple[int]]=None,
) -> np.ndarray | float :
    return gennorm.rvs(beta, loc=mu, scale=alpha, size=size, random_state=rng)

That seems to work fine and I can run an example model with the CustomDist as an observed stochastic variable just fine. However, that is not quite enough for my needs and when I implement the following model:

def fit_to_general_gaussian(t, w, dw, sample=2000, tune=1500):

    # w = observed velocity
    # dv = uncertainty on observed velocitiy
    
    # We cheat a bit in that we draw for the full sample but we subset
    # subsequently to be consistent.
    Ntotal = len(w)

    with pm.Model() as model:
        # Mean velocity 
        mu = pm.Normal("mu", mu=130, sigma=20)

        # Prior on the intrinsic velocity dispersion
        alpha = pm.Uniform("alpha", lower=0, upper=100.0)
        
        # Prior on generalised Gaussian shape
        beta = pm.Uniform("beta", lower=0, upper=4.0)

        # Draw the true velocity
        v = pm.CustomDist(
            'v',
            mu, alpha, beta, 
            logp=logp,
            random=random,
            shape=Ntotal
        )

        # Then the observed velocity 
        vobs = pm.Normal("vobs", mu=v, tau=1.0/dw**2,
                            observed=w, shape=Ntotal)

        idata = pm.sample(sample, tune=tune, return_inferencedata=True, chains=2)
    return model, idata

I get a (to me, hopefully not to others!) puzzling error:

---------------------------------------------------------------------------
KeyError                                  Traceback (most recent call last)
File ~/miniforge3/envs/pymc/lib/python3.11/site-packages/pytensor/tensor/type.py:293, in TensorType.dtype_specs(self)
    292 try:
--> 293     return self.dtype_specs_map[self.dtype]
    294 except KeyError:

KeyError: 'object'

During handling of the above exception, another exception occurred:

TypeError                                 Traceback (most recent call last)
Cell In[25], line 1

<.... Lots of possibly irrelevant error messages ....>

TypeError: Unsupported dtype for TensorType: object

I presume I have misunderstood how to use the CustomDist function - I thought it would be a straight swap-in for
v = pm.Normal("v", mu=mu, tau=1.0/sigmav**2, shape=Ntotal)
but that does not seem to be the case?

The pymc version is 5.10.0

Any suggestions most welcome!

2

gennorm.pdf doesn’t work with PyTensor TensorVariables. You have to write it using PyTensor operations or wrap gennorm in an Op.

There’s a guide focused on blackbox likelihood, but it’s the same thing for any operation you may want to define: Using a “black box” likelihood function (numpy) — PyMC example gallery

3

Thank you so much - that was indeed the fix (plus the facepalm silliness of using PDF where ln PDF was called for). In the end I rewrote it similarly to the code used for the other distributions in pymc. For reference, what I used that worked for me (I don’t know if it is useful enough to warrant a PR?):

from pytensor.tensor.variable import TensorVariable
from pytensor.tensor.random.op import RandomVariable
from typing import List, Tuple
import pytensor.tensor as pt
from pymc.distributions.continuous import Continuous
from pymc.distributions.dist_math import check_parameters
from pymc.distributions.shape_utils import rv_size_is_none

class GeneralizedGaussianRV(RandomVariable):
    name: str = "gennorm"


    ndim_supp: int = 0
    ndims_params: List[int] = [0, 0, 0]
    dtype: str = "floatX"
    _print_name: Tuple[str, str] = ("Generalized Gaussian", "\\operatorname{GenNorm}")

    def __call__(self, loc=0.0, alpha=1.0, beta=2.0, **kwargs) -> TensorVariable:
        return super().__call__(loc, alpha, beta, **kwargs)

    @classmethod
    def rng_fn(
        cls,
        rng: np.random.RandomState,
        loc: np.ndarray,
        alpha: np.ndarray,
        beta: np.ndarray,
        size: Tuple[int, ...],
    ) -> np.ndarray:
        return scipy.stats.gennorm.rvs(beta, loc=loc, scale=alpha, size=size, random_state=rng)

# Create the actual `RandomVariable` `Op`...
gennorm = GeneralizedGaussianRV()

# Subclassing `Continuous` will dispatch a default `log` transformation
class GeneralizedGaussian(Continuous):
    # This will be used by the metaclass `DistributionMeta` to dispatch the
    # class `logp` and `logcdf` methods to the `blah` `Op` defined in the last line of the code above.
    rv_op = gennorm

    # dist() is responsible for returning an instance of the rv_op.
    # We pass the standard parametrizations to super().dist
    @classmethod
    def dist(cls, mu, alpha, beta, **kwargs) -> RandomVariable:
        mu = pt.as_tensor_variable(floatX(mu))
        alpha = pt.as_tensor_variable(floatX(alpha))
        beta = pt.as_tensor_variable(floatX(beta))

        return super().dist([mu, alpha, beta], **kwargs)

    def moment(rv, size, mu, alpha, beta):
        moment, _ = pt.broadcast_arrays(mu, alpha, beta)
        if not rv_size_is_none(size):
            moment = pt.full(size, moment)
        return moment
    def logp(value, mu, alpha, beta):

        z = (value-mu)/alpha
        lp = pt.log(0.5*beta) - pt.log(2.0*alpha) - pt.gammaln(1.0/beta) - pt.pow(pt.abs(z), beta)
        
        return check_parameters(
            lp,
            alpha > 0,
            msg="alpha > 0",
        )
4

Sounds like it would be really useful for other users! We are introducing more “esoteric” distributions in pymc-experimental these days: GitHub - pymc-devs/pymc-experimental

I invite you to open a PR there!

5

Sounds good - I’ll do that, I was actually not aware of pymc-experimental, that is a very smart way to separate out development.

I thought it was merely adding documentation and tests but I can’t seem to get my previously working code working inside pymc-experimental. I presume the best place to discuss this is in the pymc-experimental Issues?

6

You can open a PR with your code, even if it’s not working. It will be easier to provide feedback that way. If you don’t want to, an issue is also fine

7

Ok! Sounds like the more productive way to proceed. PR submitted.

Thanks for the guidance!