Shape Issue with custom logp in DensityDist

Questions v5
1

Hello, I am trying to convert this notebook to PyMC v5 (from v3). I am stuck on code block 16.7. I read the updated DensityDist documentation and believe this is this should be the correct way to use that API.

def logp(value, p, majority_first):

    phi = [None] * 5

    # Probability of data
    phi[0] = pt.switch(pt.eq(value, 2), 1, 0)
    phi[1] = pt.switch(pt.eq(value, 3), 1, 0)
    phi[2] = pt.switch(pt.eq(value, 1), 1, 0)
    phi[3] = pt.ones_like(value) * 1 / 3
    phi[4] = pt.switch(
        pt.eq(majority_first, 1), pt.switch(pt.eq(value, 2), 1, 0), pt.switch(pt.eq(value, 3), 1, 0)
    )

    # Compute log ( p_s * Pr(y_i|s) )
    for i in range(5):
        phi[i] = pt.log(phi[i]) + pt.log(p[i])

    # Compute average log-probability of y_i
    return pt.sum(pm.math.logsumexp(phi, axis=0))


with pm.Model() as m16_2:
    # prior
    p = pm.Dirichlet("p", np.array([4, 4, 4, 4, 4]))

    # majority first data
    majority_first = pm.Data("majority_first", df["majority_first"].values)

    # likelihood
    y = pm.DensityDist("y", p, majority_first, logp=logp, observed=df["y"].values)

However, this code block returns this error:

ValueError: Could not broadcast dimensions. Incompatible shapes were [(ScalarConstant(ScalarType(int64), data=5),), (ScalarConstant(ScalarType(int64), data=629),)].

It should be noted that the shape of df["y"].values is (629,)

Last updated: Tue Jan 09 2024

Python implementation: CPython
Python version       : 3.11.7
IPython version      : 8.19.0

numpy     : 1.26.3
pytensor  : 2.18.4
arviz     : 0.17.0
seaborn   : 0.12.2
pymc      : 5.10.3
pandas    : 2.1.4
matplotlib: 3.8.2

Watermark: 2.4.3

Thank you in advance for any direction on this!

2

You may need to specify ndims_params=[1, 0] to tell PyMC that the probability of a single observation is parametrized by a vector as the first parameter and a scalar for the second (if I got that correctly)

There are more details about this in this guide: Distribution Dimensionality — PyMC 5.10.3 documentation

Particularly this section: Distribution Dimensionality — PyMC 5.10.3 documentation

It works the same for CustomDist. By default PyMC assumes each observation is associated with scalar parameters

3

Thanks, Ricardo. This helped get the model to compile. When I sample, however, I get this surprising error message:

<<!! BUG IN FGRAPH.REPLACE OR A LISTENER !!>> <class 'TypeError'> The type of the replacement (Matrix(float64, shape=(1, 1))) must be compatible with the type of the original Variable (Matrix(float64, shape=(1, 629))). local_add_neg_to_sub
ERROR (pytensor.graph.rewriting.basic): Rewrite failure due to: local_add_neg_to_sub
ERROR (pytensor.graph.rewriting.basic): node: Add([[-1.09861 ... 09861229]], ExpandDims{axes=[0, 1]}.0)
ERROR (pytensor.graph.rewriting.basic): TRACEBACK:
ERROR (pytensor.graph.rewriting.basic): Traceback (most recent call last):
  File "c:\Users\mason\AppData\Local\Continuum\anaconda3\envs\stat-rethink2-pymc\Lib\site-packages\pytensor\graph\rewriting\basic.py", line 1968, in process_node
    fgraph.replace_all_validate_remove(  # type: ignore
  File "c:\Users\mason\AppData\Local\Continuum\anaconda3\envs\stat-rethink2-pymc\Lib\site-packages\pytensor\graph\features.py", line 626, in replace_all_validate_remove
    chk = fgraph.replace_all_validate(replacements, reason=reason, **kwargs)
          ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "c:\Users\mason\AppData\Local\Continuum\anaconda3\envs\stat-rethink2-pymc\Lib\site-packages\pytensor\graph\features.py", line 571, in replace_all_validate
    fgraph.replace(r, new_r, reason=reason, verbose=False, **kwargs)
  File "c:\Users\mason\AppData\Local\Continuum\anaconda3\envs\stat-rethink2-pymc\Lib\site-packages\pytensor\graph\fg.py", line 508, in replace
    self.change_node_input(
  File "c:\Users\mason\AppData\Local\Continuum\anaconda3\envs\stat-rethink2-pymc\Lib\site-packages\pytensor\graph\fg.py", line 428, in change_node_input
    raise TypeError(
TypeError: The type of the replacement (Matrix(float64, shape=(1, 1))) must be compatible with the type of the original Variable (Matrix(float64, shape=(1, 629))).

However, the model appears to sample just fine, and I am able to replicate the summary outputs from the previous model version. Is this perhaps an error in model specification or did this uncover a true bug?

4

Could be either. Any chance you could get me a fully reproducible snippet? The shorter the better but anything should do

5

Sure thing:

import arviz as az
import numpy as np
import pandas as pd
import pymc as pm
from pytensor import tensor as pt
from numpy.random import default_rng

RANDOM_SEED = 1234
np.random.seed(RANDOM_SEED)
rng = default_rng(RANDOM_SEED)
az.style.use("arviz-darkgrid")
az.rcParams["stats.hdi_prob"] = 0.89

df = pd.read_csv("https://raw.githubusercontent.com/rmcelreath/rethinking/master/data/Boxes.csv", delimiter=";")

def logp(value, p, majority_first):

    phi = [None] * 5

    # Probability of data
    phi[0] = pt.switch(pt.eq(value, 2), 1, 0)
    phi[1] = pt.switch(pt.eq(value, 3), 1, 0)
    phi[2] = pt.switch(pt.eq(value, 1), 1, 0)
    phi[3] = pt.ones_like(value) * 1 / 3
    phi[4] = pt.switch(
        pt.eq(majority_first, 1), pt.switch(pt.eq(value, 2), 1, 0), pt.switch(pt.eq(value, 3), 1, 0)
    )

    # Compute log ( p_s * Pr(y_i|s) )
    for i in range(5):
        phi[i] = pt.log(phi[i]) + pt.log(p[i])

    # Compute average log-probability of y_i
    return pt.sum(pm.math.logsumexp(phi, axis=0))


with pm.Model() as m16_2:
    # prior
    p = pm.Dirichlet("p", np.array([4, 4, 4, 4, 4]))

    # majority first data
    majority_first = pm.ConstantData("majority_first", df["majority_first"].values)

    # likelihood
    y = pm.DensityDist(
        "y", p, majority_first, logp=logp, ndims_params=[1, 0], observed=df["y"].values
    )

    idata = pm.sample()
6

I can confirm that’s a bug. Now tracked at: `local_add_neg_to_sub` rewrite gives wrong results with negative constants · Issue #584 · pymc-devs/pytensor · GitHub

It shouldn’t affect the correctness of your code, because PyTensor is correctly rejecting the rewrite suggestion (while also being verbose about having found a bug in the codebase, not in your graph)

7

Small suggestion, I would perhaps replace that loop by:

    # Compute log ( p_s * Pr(y_i|s) )
    phi = pt.concatenate([p[None] for p in phi] , axis=0)  # PyTensor doesn't have `pt.vstack`
    phi = pt.log(phi) + pt.log(p[:, None])
8

Thanks for these suggestions. Both work and there is no longer any bug messages when I call pm.sample. I appreciate your insights.

1 Like