JAGS issue (non-conforming error by inprod)

Question

I am relatively new to using JAGS, which I am using rjags package in R to run my coding scheme. Let me summarize as best as I can. I am implementing a JAGS model based off of the coding scheme located in the article "Sensitivity analyses for the principal ignorability assumption using multiple imputation by Craig Wang, Yufen Zhang, Fabrizia Mealli, Björn Bornkamp 2022-01-10". The idea is to use a joint imputation model using the principal ignorability assumption for principal stratum. It comprises of a joint likelihood of P(S|X) and P(Y|S,X), which can be shown with the JAGS model and setup here:

set.seed(06262023)

N = 335
TRT = sample(c("treatment", "control"), size=N, replace=TRUE, prob=c(0.50, 0.50))
X1 = sample(18:90, size=N, replace=TRUE)
X2 = sample(c("TRUE", "FALSE"), size=N, replace=TRUE, prob=c(0.35, 0.65))
X3 = sample(c("TRUE", "FALSE"), size=N, replace=TRUE, prob=c(0.40, 0.60))
X4 = sample(c("SECONDARY RESISTANT", "SENSITIVE"), size=N, replace=TRUE, prob=c(0.70, 0.30))
X5 = sample(c("<3", ">=3"), size=N, replace=TRUE, prob=c(0.80, 0.20))
PFS = rexp(N, exp(-2))
EVENT = rbinom(N, 1, 0.3)
dAE = rbinom(N, 1, 0.2)

dat = data.frame(TRT, X1, X2, X3, X4, X5, PFS, EVENT, dAE)

modstr <- "model{
  for(i in 1:N){
    ## survival model for Y(0)
    mu[i] <- inprod(beta_s[], X_s[i,])
    I[i] ~ dinterval(Y0[i], cens[i])
    
    ## binary model for S(1) (trt disc due to AE on trt)
    logit(p[i]) <- inprod(beta_b[], X_b[i,])
    S1[i] ~ dbern(p[i])
    lambda[i] <- exp(mu[i] + gamma*S1[i])
    Y0[i] ~ dweib(alpha, lambda[i])
  }
 
  alpha ~ dgamma(0.1, 0.1)
  for(b in 1:7){
    beta_s[b] ~ dnorm(0, 0.01)
    beta_b[b] ~ dnorm(0, 0.01)
  }
}
"
cat(modstr, file="mod.txt")

dat <- dat %>% 
  mutate(I = ifelse(TRT == "treatment", 1, !EVENT)) %>% 
  mutate(Y0 = ifelse(TRT == "treatment", NA, ifelse(I == 1, NA, PFS))) %>% 
  mutate(cens = ifelse(TRT == "treatment", 0, ifelse(I == 1, PFS, PFS + 1))) %>% 
  mutate(S1 = ifelse(TRT == "treatment", dAE, NA)) %>% filter(X4 != "") %>% droplevels()

X <- model.matrix(~ X1 + X2 + X3 + X4 + X5, data = dat)

init.beta_b <- glm(dAE ~ X1 + X2 + X3 + X4 + X5, 
                   data = dat %>% filter(TRT == "treatment"), 
                   family = "binomial")$coefficient

mod.beta_s <- survreg(Surv(PFS, EVENT) ~ S1 + X1 + X2 + X3 + X4 + X5, 
                      dist = 'weibull', data = dat %>% filter(TRT == "treatment"))
init.beta_s <- -mod.beta_s$coefficients[-2]/mod.beta_s$scale

gamma_0 = 0.5
dat.jags <- list(I=dat$I, cens=dat$cens, Y0=dat$Y0, S1=dat$S1, N=nrow(dat), X_b = X, X_s = X, gamma = gamma_0)
inits <- function() list(Y0 = Y02, alpha = 1, beta_s = init.beta_s, beta_b = init.beta_b,
                         .RNG.name = "base::Wichmann-Hill", .RNG.seed = seed) 
parSave <- c("Y0", "S1", "alpha", "beta_s", "beta_b") 

Y02 <- dat$Y0
ind <- is.na(dat$Y0)
Y02[!ind] <- NA
Y02[ind] <- rexp(sum(ind), 1/median(dat$Y0, na.rm = T))

mod <- jags.model("mod.txt", dat.jags, inits, n.chains = 2, n.adapt = 100, quiet = T)
sims <- coda.samples(mod, parSave, n.iter = 10000, thin = 5)

S1_ind <- grep("S1", colnames(sims[[1]]))
nSim <- nrow(sims[[1]])
hr1s <- hr0s <- hr1_vars <- hr0_vars <- numeric(nSim)
for (i in 1:nSim) {
  dat_imp <- cbind(dat, sims[i, S1_ind])
  names(dat_imp)[ncol(dat_imp)] <- "S1_imp"
  fit1 <- coxph(Surv(PFS, EVENT) ~ TRT+ X1 + X2 + X3 + X4 + X5,
                data = dat_imp %>% filter(S1_imp == 1))
  fit0 <- coxph(Surv(PFS, EVENT) ~ TRT+ X1 + X2 + X3 + X4 + X5,
                data = dat_imp %>% filter(S1_imp == 0))
  hr1s[i] <- coef(fit1)[1]
  hr0s[i] <- coef(fit0)[1]
  hr1_vars[i] <- vcov(fit1)[1, 1]
  hr0_vars[i] <- vcov(fit0)[1, 1]
}

out <- data.frame(S1_1_est = c(mean(hr1s)), S1_1_var = c(mean(hr1_vars) + var(hr1s)), 
                  S1_0_est = c(mean(hr0s)), S1_0_var = c(mean(hr0_vars) + var(hr0s)))
est <- data.frame(coef.s0 = exp(out$S1_0_est),
                  coef.s0.h = exp(out$S1_0_est - 1.96*sqrt(out$S1_0_var)),
                  coef.s0.l = exp(out$S1_0_est + 1.96*sqrt(out$S1_0_var)),
                  coef.s1 = exp(out$S1_1_est),
                  coef.s1.h = exp(out$S1_1_est - 1.96*sqrt(out$S1_1_var)),
                  coef.s1.l = exp(out$S1_1_est + 1.96*sqrt(out$S1_1_var)))

I am getting a runtime error such as: Error in jags.model("mod.txt", dat.jags, inits, n.chains = 2, n.adapt = 100,: RUNTIME ERROR: Non-conforming parameters in function inprod

I think it is coming within my jags model about non-conforming beta and X matrix, but I don't know why they are not conforming. It seems to work just fine outside of the algorithm but not sure what is going on now. If anyone can help, that would be great! Thank you!