Sampling many dimensions, with different PDF in each dimension

Hi Ricardo,

It took me some time to correctly implement pymc5 to make CustomDist work. However, once implemented, the script does not work. I don’t make it to understand the output message . Here the info. Thank you very much! (Maybe there is another option simpler than CustomDist?)

Best, Ezequiel.

def diag_mv_exponential_dist(lambda1, lambda2, size):
    if size is None:
        size = ()
    return pm.Exponential.dist([lambda1, lambda2], shape=tuple(size) + (2,))

def diag_mv_exponential_logp(value, lambda1, lambda2):
    # Contract logp of the pairs
    return pm.logp(pm.Exponential.dist([lambda1, lambda2]), value).sum(-1)

with pm.Model() as m:
    w = pm.Dirichlet('w', a=np.array([1, 1]))
    mu0 = pm.Normal("mu0", mu=120, sigma=15)
    mu1 = pm.Normal("mu1", mu=124, sigma=2)
    sigma0 = pm.HalfNormal("sigma0", sigma=15)
    sigma1 = pm.HalfNormal("sigma1", sigma=1)        
    lambda0 = pm.Normal("lambda0", mu=0.05, sigma=0.02)
    lambda1 = pm.Normal("lambda1", mu=0.5, sigma=0.2)
    
    diag_mv_exponential = pm.CustomDist.dist(lambda0, lambda1, dist=diag_mv_exponential_dist, logp=diag_mv_exponential_logp, ndim_supp=1, class_name="diag_mv_exponential")
    diag_mv_normal = pm.MvNormal.dist(mu=[mu0,mu1], cov=[[sigma0,0],[0,sigma1]])
    mymixture = [diag_mv_normal, diag_mv_exponential]

    Y_obs = pm.Mixture("Y_obs", w = w, comp_dists=mymixture, observed=data)

And sends me the following message, which starts with a " UserWarning: CustomDist with dist function is still experimental. Expect bugs!"

/home/sequi/miniconda3/envs/pymc_env/lib/python3.11/site-packages/pymc/distributions/distribution.py:623: UserWarning: CustomDist with dist function is still experimental. Expect bugs!
  warnings.warn(
---------------------------------------------------------------------------
AttributeError                            Traceback (most recent call last)
File ~/miniconda3/envs/pymc_env/lib/python3.11/site-packages/pymc/pytensorf.py:436, in floatX(X)
    435 try:
--> 436     return X.astype(pytensor.config.floatX)
    437 except AttributeError:
    438     # Scalar passed

AttributeError: 'list' object has no attribute 'astype'

During handling of the above exception, another exception occurred:

TypeError                                 Traceback (most recent call last)
TypeError: float() argument must be a string or a real number, not 'TensorVariable'

The above exception was the direct cause of the following exception:

ValueError                                Traceback (most recent call last)
Cell In[50], line 19
     16 lambda0 = pm.Normal("lambda0", mu=0.05, sigma=0.02)
     17 lambda1 = pm.Normal("lambda1", mu=0.5, sigma=0.2)
---> 19 diag_mv_exponential = pm.CustomDist.dist(lambda0, lambda1, dist=diag_mv_exponential_dist, logp=diag_mv_exponential_logp, ndim_supp=1, class_name="diag_mv_exponential")
     20 diag_mv_normal = pm.MvNormal.dist(mu=[mu0,mu1], cov=[[sigma0,0],[0,sigma1]])
     21 mymixture = [diag_mv_normal, diag_mv_exponential]

File ~/miniconda3/envs/pymc_env/lib/python3.11/site-packages/pymc/distributions/distribution.py:1023, in CustomDist.dist(cls, class_name, dist, random, logp, logcdf, moment, ndim_supp, ndims_params, dtype, *dist_params, **kwargs)
   1021 cls.check_valid_dist_random(dist, random, dist_params)
   1022 if dist is not None:
-> 1023     return _CustomSymbolicDist.dist(
   1024         *dist_params,
   1025         class_name=class_name,
   1026         dist=dist,
   1027         logp=logp,
   1028         logcdf=logcdf,
   1029         moment=moment,
   1030         ndim_supp=ndim_supp,
   1031         **kwargs,
   1032     )
   1033 else:
   1034     return _CustomDist.dist(
   1035         *dist_params,
   1036         class_name=class_name,
   (...)
   1044         **kwargs,
   1045     )

File ~/miniconda3/envs/pymc_env/lib/python3.11/site-packages/pymc/distributions/distribution.py:641, in _CustomSymbolicDist.dist(cls, class_name, dist, logp, logcdf, moment, ndim_supp, dtype, *dist_params, **kwargs)
    633 if moment is None:
    634     moment = functools.partial(
    635         default_moment,
    636         rv_name=class_name,
    637         has_fallback=True,
    638         ndim_supp=ndim_supp,
    639     )
--> 641 return super().dist(
    642     dist_params,
    643     class_name=class_name,
    644     logp=logp,
    645     logcdf=logcdf,
    646     dist=dist,
    647     moment=moment,
    648     ndim_supp=ndim_supp,
    649     **kwargs,
    650 )

File ~/miniconda3/envs/pymc_env/lib/python3.11/site-packages/pymc/distributions/distribution.py:387, in Distribution.dist(cls, dist_params, shape, **kwargs)
    385 ndim_supp = getattr(cls.rv_op, "ndim_supp", None)
    386 if ndim_supp is None:
--> 387     ndim_supp = cls.rv_op(*dist_params, **kwargs).owner.op.ndim_supp
    388 create_size = find_size(shape=shape, size=size, ndim_supp=ndim_supp)
    389 rv_out = cls.rv_op(*dist_params, size=create_size, **kwargs)

File ~/miniconda3/envs/pymc_env/lib/python3.11/site-packages/pymc/distributions/distribution.py:670, in _CustomSymbolicDist.rv_op(cls, class_name, dist, logp, logcdf, moment, size, ndim_supp, *dist_params)
    666 dummy_dist_params = [dist_param.type() for dist_param in dist_params]
    667 with BlockModelAccess(
    668     error_msg_on_access="Model variables cannot be created in the dist function. Use the `.dist` API"
    669 ):
--> 670     dummy_rv = dist(*dummy_dist_params, dummy_size_param)
    671 dummy_params = [dummy_size_param] + dummy_dist_params
    672 dummy_updates_dict = collect_default_updates(dummy_params, (dummy_rv,))

Cell In[50], line 4, in diag_mv_exponential_dist(lambda1, lambda2, size)
      2 if size is None:
      3     size = ()
----> 4 return pm.Exponential.dist([lambda1, lambda2], shape=tuple(size) + (2,))

File ~/miniconda3/envs/pymc_env/lib/python3.11/site-packages/pymc/distributions/continuous.py:1314, in Exponential.dist(cls, lam, *args, **kwargs)
   1312 @classmethod
   1313 def dist(cls, lam, *args, **kwargs):
-> 1314     lam = at.as_tensor_variable(floatX(lam))
   1316     # PyTensor exponential op is parametrized in terms of mu (1/lam)
   1317     return super().dist([at.reciprocal(lam)], **kwargs)

File ~/miniconda3/envs/pymc_env/lib/python3.11/site-packages/pymc/pytensorf.py:439, in floatX(X)
    436     return X.astype(pytensor.config.floatX)
    437 except AttributeError:
    438     # Scalar passed
--> 439     return np.asarray(X, dtype=pytensor.config.floatX)

ValueError: setting an array element with a sequence.