Using CustomDist as a stochastic variable - causing TensorType error

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(
            mu, alpha, beta, 

        # 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/, 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!

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

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)

    def rng_fn(
        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
    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(
            alpha > 0,
            msg="alpha > 0",

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!

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?

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

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

Thanks for the guidance!