I think I made it work. See I’m using `pm.CustomDist`

instead of `pm.DensityDist`

. I had a look at the PyMC codebase and it looks like DensityDist will be deprecated at some point and CustomDist is the new thing.

```
import arviz as az
import bambi as bmb
import numpy as np
import pandas as pd
import pymc as pm
import pytensor.tensor as pt
from bambi.families.univariate import UnivariateFamily
rng = np.random.default_rng(1234)
data = pd.DataFrame(
{
"score": np.concatenate([rng.normal(size=1000), rng.normal(loc=1, size=1000)]),
"weights": np.repeat([1, 2], 1000)
}
)
def logp(value, mu, sigma, weights):
# This is the same as the logp of a normal, but we have the weights multiplying
return weights * pm.Normal.logp(value, mu, sigma)
def random(mu, sigma, rng=None, size=None):
# Weights don't make sense when generating new observations.
# They are a property of actual observations
return rng.normal(loc=mu, scale=sigma, size=size)
class WeightedNormal:
def __new__(self, name, mu, sigma, weights, **kwargs):
return pm.CustomDist(name, mu, sigma, weights, logp=logp, **kwargs)
@classmethod
def dist(cls, mu, sigma, **kwargs):
return pm.CustomDist.dist(
mu, sigma, class_name="WeightedNormal", random=random, **kwargs
)
# Create custom family.
# We need to use the weights in a different way, so we need 'transform_backend_kwargs'.
class WeightedNormalFamily(UnivariateFamily):
SUPPORTED_LINKS = {"mu": ["identity", "log"], "sigma": ["log"]}
@staticmethod
def transform_backend_kwargs(kwargs):
if "observed" in kwargs:
observed = kwargs["observed"]
kwargs["observed"] = observed[:, 0]
kwargs["weights"] = observed[:, 1]
return kwargs
# Create the custom family instance
likelihood = bmb.Likelihood("WeightedNormal", params=["mu", "sigma"], parent="mu", dist=WeightedNormal)
links = {"mu": "identity", "sigma": "log"}
wn_family = WeightedNormalFamily("weighted-normal", likelihood, links)
# Create the model in Bambi using the custom family
priors = {
"mu": bmb.Prior("Normal", mu=0, sigma=1),
"sigma": bmb.Prior("HalfNormal", sigma=1),
}
model = bmb.Model("c(score, weights) ~ 1", data, family=wn_family, priors=priors)
model.build()
model.graph()
idata = model.fit()
az.summary(idata)
idata_2 = model.predict(idata, data = data, kind="pps", inplace = False)
```

See it’s a lot of code. Hopefully, we can have this inside Bambi in the future.