Dynamic shaping, "round" function, JAX, and a "few" more questions

Questions v5
, ,
17

Hi @ricardoV94, I updated my PyMC version to the latest version 5.6, and tested whether it was now possible to take a CustomDist of another CustomDist. Unfortunately, I get the same error as before:

BlockModelAccessError: Model variables cannot be created in the dist function. Use the .dist API

In the case of my model, the first CustomDist is created by providing the custom functions random, logp, and logcdf. Below is an example of a simplified and reproducible code where this first CustomDist “_cstm_dist1” is replaced by a “dummy” CustomDist used to simulate this distribution by wrapping a common Normal dist. Although not exactly equivalent to my model, the same error is obtained, please see below. The error occurs at the line _cstm_dist = _cstm_dist2(mu, sigma).

import numpy as np
import pymc as pm
import pytensor.tensor as pt

#**"_cstm_dist3" definition
def _rnded_norm_dist(mu, sigma, size=None):
    _dist = pt.round(pm.Normal.dist(mu=mu, sigma=sigma, shape=size))
    return _dist

def _cstm_dist3(mu, sigma, size=None):
    _dist = pm.CustomDist.dist(mu, sigma,
                dist=_rnded_norm_dist,
                shape=size)
    return _dist

#**"_cstm_dist2" definition
def _norm_dist(mu, sigma, size=None):
    _dist = pm.Normal.dist(mu=mu, sigma=sigma, shape=size)
    return _dist

def _cstm_dist1(mu, sigma, size=None):
    _dist = pm.CustomDist.dist(mu, sigma, dist=_norm_dist, shape=size)
    return _dist

def _rnded_cstm_dist2(mu, sigma, size=None):
    _dist = pt.round(_cstm_dist1(mu, sigma, size=size))
    return _dist

def _cstm_dist2(mu, sigma, size=None):
    _dist = pm.CustomDist.dist(mu, sigma, dist=_rnded_cstm_dist2, shape=size)
    return _dist

mu = 0.
sigma = 1.
data = np.array([0.5, 0.5+0.0001])

#**test#1 - "_cstm_dist3"
# WORKS!
_cstm_dist = _cstm_dist3(mu, sigma)
vv = pt.vector(name="v")
_cstm_dist_logp = pm.logp(rv=_cstm_dist, value=vv)
rr3 = np.exp(_cstm_dist_logp.eval({vv: data}))
print(rr3)

#**test#2 - "_cstm_dist2"
# ERROR RAISED: "BlockModelAccessError: Model variables cannot be created
#               "in the dist function. Use the `.dist` API"
_cstm_dist = _cstm_dist2(mu, sigma)
vv = pt.vector(name="v")
_cstm_dist_logp = pm.logp(rv=_cstm_dist, value=vv)
rr2 = np.exp(_cstm_dist_logp.eval({vv: data}))
print(rr2)

I am sorry it’s not the outcome we expected … there may also be a mistake in the logic of my implementation?

Thank you for your time and help.