Numpyro render_model failure?

[fehiepsi]

@jecampagne

def diag_mean_fn(x, params):
    R0 = params["R0"]
    v  = params["v"]
    k  = params["k"]
    tau =  params["tau"]
    return jnp.piecewise(
        x, [x < 0, x >= 0],
        [lambda x: R0 + v*x, 
         lambda x: R0 + v*x - k*(1.-jnp.exp(-x/tau))
        ])

I think the current provenance approach does not work with jax control flow. I think you can use

return jnp.where(x < 0, R0 + v*x, R0 + v*x - k*(1.-jnp.exp(-x/tau)))

instead. Here is a reference for plate: Plate notation - Wikipedia