Custom distributions in PyMc4

Questions v5
1

I am trying to create a custom distribution (generalized normal) following the explanation in the docs. I wrote the following code:

import numpy as np
import aesara.tensor as at

from typing import List, Tuple

from aesara.tensor.var import TensorVariable
from aesara.tensor.random.op import RandomVariable
from pymc.aesaraf import floatX, intX
from pymc.distributions.distribution import Continuous
from pymc.distributions.dist_math import check_parameters


class GenNormRV(RandomVariable):
    name: str = "GenNorm"
    ndim_supp: int = 0
    ndims_params: List[int] = [0, 0, 0]
    dtype: str = "floatX"
    _print_name: Tuple[str, str] = ("GenNorm", "GGD")

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


class GenNorm(Continuous):
    rv_op = GenNormRV

    @classmethod
    def dist(cls, beta, loc, scale, *args, **kwargs):
        beta = at.as_tensor_variable(floatX(beta))
        loc = at.as_tensor_variable(floatX(loc))
        scale = at.as_tensor_variable(floatX(scale))
        return super().dist([beta, loc, scale], *args, **kwargs)

    def moment(rv, size, beta, loc, scale):
        moment, _ = at.broadcast_arrays(beta, loc, scale)
        if not rv_size_is_none(size):
            moment = at.full(size, moment)
        return moment

    def logp(value, beta, loc, scale):
        return check_parameters(
            at.log(beta / (2 * scale)) - at.gammaln(1.0 / beta) -
            (at.abs_(value - loc) / scale)**beta, beta >= 0, scale >= 0)

    def logcdf(value, beta, loc, scale):
        b = value - loc
        c = 0.5 * b / at.abs_(b)
        return (0.5 + c) - c * at.gammaincc(1.0 / beta,
                                            at.abs_(b / scale)**beta)

When I am trying out this example using the code from the docs:

import pymc as pm
from pymc.distributions.distribution import moment

d = GenNorm.dist(beta=1, loc=0, scale=1)

I get the following (confusing) error:

---------------------------------------------------------------------------
TypeError                                 Traceback (most recent call last)
Input In [10], in <cell line: 4>()
      1 import pymc as pm
      2 from pymc.distributions.distribution import moment
----> 4 d = GenNorm.dist(beta=1, loc=0, scale=1)
      6 # Test that the returned blah_op is still working fine
      7 blah.eval()

Input In [9], in GenNorm.dist(cls, beta, loc, scale, *args, **kwargs)
     13 loc = at.as_tensor_variable(floatX(loc))
     14 scale = at.as_tensor_variable(floatX(scale))
---> 15 return super().dist([beta, loc, scale], *args, **kwargs)

File ~/repos/magisterska/.venv/lib/python3.10/site-packages/pymc/distributions/distribution.py:351, in Distribution.dist(cls, dist_params, shape, **kwargs)
    346 create_size, ndim_expected, ndim_batch, ndim_supp = find_size(
    347     shape=shape, size=size, ndim_supp=cls.rv_op.ndim_supp
    348 )
    349 # Create the RV with a `size` right away.
    350 # This is not necessarily the final result.
--> 351 rv_out = cls.rv_op(*dist_params, size=create_size, **kwargs)
    353 # Replicate dimensions may be prepended via a shape with Ellipsis as the last element:
    354 if shape is not None and Ellipsis in shape:

TypeError: RandomVariable.__init__() got an unexpected keyword argument 'size'

As far as I can see in the source code, other distributions in distributions/continuous.py have been created in the similar way. Can you tell me what I am doing wrong?

2

Actually, I found out that this was the bug - this was supposed to be

rv_op = GenNormRV()
3 Likes