Hi all,

I have a model in which in which subjects can get one of two treatments. The distribution of the response depends on the treatment; treatment A results in a negative HalfNormal distributed response, and treatment B results in a positive TruncatedNormal response. In my model, I created a switch statement, and applied the correct model based on the treatment. The problem I have is that I cannot treat the output of the switch step as an observed value. In the model below, is there a way I can use the output of the switch statement as the observed data?

model = pm.Model()

with model:

```
trt1_idx = (df['treatment']==1).astype(int).values
trt2_idx = (df['treatment']==2).astype(int).values
obs_data=df['response'].values
#Treatment 1
mu_age = pm.Normal('mu_BLd', 20, 3)
Bage = pm.Normal('Bage', mu_age)
mu_wt = pm.Normal('mu_wt', 150, 10)
Bwt = pm.Normal('Bwt' , mu_wt)
rsp1_mu = pm.Deterministic('rsp1_mu', (df['age'].values**mu_age ) * (df['weight'].values**Bwt) )
rsp1 = -1*pm.HalfNormal('rsp1', rsp1_mu)
# Treatment 2
mu_ht = pm.Normal('mu_ht', 80, 5)
Bht = pm.Normal('Bht' , mu_ht)
rsp2_mu = pm.Deterministic('rsp2_mu', (df['height'].values**Bht) )
rsp2 = pm.TruncatedNormal('rsp2', rsp2_mu, lower=0)
#Combined Response
response=pm.math.switch(df['treatment']==1, rsp1[trt1_idx], rsp2[trt2_idx], shape = df.shape[0], observed=obs_data)
```