Min of two samples inside model with NUTS

I’m hoping this is a pretty basic question. I have a model where I want to constrain the sum of two parameters to be negative (one is deterministic given from the input and one from sample) and I’m wondering what is the best way to accomplish this. Without the constraint the MCMC process runs smoothly but NUTS seems to have problems when I use jnp.mean

For example:

def model(df):
    ...
    with numpyro.plate("sc_i", K):
        mu_1 = numpyro.sample('mu_1', dist.Normal(df.var.values,1))

    with numpyro.plate("data", N):
        beta = numpyro.deterministic('beta', jnp.min(df.var.values + mu_1[k],0))
        p = expit(alpha + beta*var)
        numpyro.sample('probs', dist.Delta(p), obs=p)
        numpyro.sample("obs", dist.Binomial(x, probs = p), obs = n)

Essentially, mu_1[k] is an offset from the precomputed values var, but I still need to hold the negativity constraint. I’m also wondering if this might be easier to do post model fit using conditioning. Any thoughts?

Solved my own problem by feeding the var vaules into the high parameter of the truncated normal distribution. See below for anyone who stumbles across a similar problem.

def model(df):
    ...
    with numpyro.plate("sc_i", K):
        mu_1 = numpyro.sample('mu_1', dist.TruncatedNormal(df.var.values,1,high=-df.var.values))

    with numpyro.plate("data", N):
        p = expit(alpha + mu_1[k]*var)
        numpyro.sample('probs', dist.Delta(p), obs=p)
        numpyro.sample("obs", dist.Binomial(x, probs = p), obs = n)