# Using CustomDist as a stochastic variable - causing TensorType error

**URL:** <https://discourse.pymc.io/t/using-customdist-as-a-stochastic-variable-causing-tensortype-error/13436>\
**Category:** v5\
**Tags:** pytensor\
**Created:** [December 8, 2023, 3:27pm UTC](https://discourse.pymc.io/t/using-customdist-as-a-stochastic-variable-causing-tensortype-error/13436 "2023-12-08T15:27:45Z")\
**Posts on this page:** 7\
**Page:** 1

<div class="post-metadata">

**Author:** ![jarle](https://avatars.discourse-cdn.com/v4/letter/j/cab0a1/32.png) [@jarle](https://discourse.pymc.io/u/jarle)\
**Post date:** [December 8, 2023, 3:27pm UTC](https://discourse.pymc.io/t/using-customdist-as-a-stochastic-variable-causing-tensortype-error/13436/1 "2023-12-08T15:27:46Z")

</div>

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:

```auto
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:

```auto
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:

```auto
---------------------------------------------------------------------------
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!

---

<div class="post-metadata">

**Author:** ![ricardoV94](https://yyz2.discourse-cdn.com/flex036/user_avatar/discourse.pymc.io/ricardov94/32/5775_2.png) [@ricardoV94](https://discourse.pymc.io/u/ricardoV94)\
**Post date:** [December 8, 2023, 9:26pm UTC](https://discourse.pymc.io/t/using-customdist-as-a-stochastic-variable-causing-tensortype-error/13436/2 "2023-12-08T21:26:05Z")

</div>

`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](https://www.pymc.io/projects/examples/en/latest/howto/blackbox_external_likelihood_numpy.html)

---

<div class="post-metadata">

**Author:** ![jarle](https://avatars.discourse-cdn.com/v4/letter/j/cab0a1/32.png) [@jarle](https://discourse.pymc.io/u/jarle)\
**Post date:** [December 8, 2023, 10:19pm UTC](https://discourse.pymc.io/t/using-customdist-as-a-stochastic-variable-causing-tensortype-error/13436/3 "2023-12-08T22:19:59Z")

</div>

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?):

```auto
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",
        )

```

---

<div class="post-metadata">

**Author:** ![ricardoV94](https://yyz2.discourse-cdn.com/flex036/user_avatar/discourse.pymc.io/ricardov94/32/5775_2.png) [@ricardoV94](https://discourse.pymc.io/u/ricardoV94)\
**Post date:** [December 9, 2023, 11:24am UTC](https://discourse.pymc.io/t/using-customdist-as-a-stochastic-variable-causing-tensortype-error/13436/4 "2023-12-09T11:24:31Z")

</div>

> [@jarle](#):
>
> I don’t know if it is useful enough to warrant a PR?

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](https://github.com/pymc-devs/pymc-experimental)

I invite you to open a PR there!

---

<div class="post-metadata">

**Author:** ![jarle](https://avatars.discourse-cdn.com/v4/letter/j/cab0a1/32.png) [@jarle](https://discourse.pymc.io/u/jarle)\
**Post date:** [December 10, 2023, 7:28pm UTC](https://discourse.pymc.io/t/using-customdist-as-a-stochastic-variable-causing-tensortype-error/13436/5 "2023-12-10T19:28:55Z")

</div>

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?

---

<div class="post-metadata">

**Author:** ![ricardoV94](https://yyz2.discourse-cdn.com/flex036/user_avatar/discourse.pymc.io/ricardov94/32/5775_2.png) [@ricardoV94](https://discourse.pymc.io/u/ricardoV94)\
**Post date:** [December 10, 2023, 8:32pm UTC](https://discourse.pymc.io/t/using-customdist-as-a-stochastic-variable-causing-tensortype-error/13436/6 "2023-12-10T20:32:58Z")

</div>

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

---

<div class="post-metadata">

**Author:** ![jarle](https://avatars.discourse-cdn.com/v4/letter/j/cab0a1/32.png) [@jarle](https://discourse.pymc.io/u/jarle)\
**Post date:** [December 11, 2023, 12:27am UTC](https://discourse.pymc.io/t/using-customdist-as-a-stochastic-variable-causing-tensortype-error/13436/7 "2023-12-11T00:27:20Z")

</div>

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

Thanks for the guidance!