Power prior with CustomDist


I am trying to implement the power prior by Ibrahim & Chen (2015). The model is the following:

Initial prior on the parameter \theta: p(\theta)
Source study data: y_S \sim \mathcal{N}(\theta, \sigma_S^2) (\sigma_S^2 known)
Target study data: y_T \sim \mathcal{N}(\theta, \sigma_T^2) (\sigma_T^2 known)
Power prior on theta for analyzing the new study: \theta | y_S \propto p(y_S|\theta)^\gamma p(\theta), where the power parameter \gamma is fixed, and used to discount source study data in the context of Bayesian borrowing.

I’m trying to implement this model using PyMC 5.0.2 with pm.CustomDist.

import numpy as np
import pymc as pm
from pytensor.tensor import TensorVariable

source_data = np.random.randn(100) + 1
target_data = np.random.randn(10) - 1

def logp(value: TensorVariable, theta_s: TensorVariable, source_data: TensorVariable, y_s: TensorVariable, gamma: TensorVariable) -> TensorVariable:
    log_initial_prior = pm.logp(theta_s, value) 
    log_source_likelihood = pm.logp(y_s,source_data).sum()
    log_unnormalized_density = log_initial_prior + gamma*log_source_likelihood
    return log_unnormalized_density 

# Power prior parameter
gamma = 0.5

with pm.Model() as model:
    # Source data parameter
    theta_s = pm.Normal("theta_s", mu= 0, sigma=1) 
    # Source study observations
    y_s = pm.Normal("y_s", mu=theta_s, sigma=1, observed = source_data) 

    # Power prior on theta:
    theta_t = pm.CustomDist("theta_t", theta_s, source_data, y_s, gamma, logp = logp)

    # Target study obeservations
    y_t =  pm.Normal("y_t", mu=theta_t, sigma=1, observed= target_data) 

    idata = pm.sample(1000)

However, I get the following error:

AttributeError: 'NoneType' object has no attribute 'op'.

Side note: my end goal is to extend this model to handle unknown power parameter \gamma (aka the “modified power prior”).
Thanks a lot for your help.