How to write an arbitrary number of models (functions) in NumPyro?

Question

I have 5 NumPyro models of the form:

def patient_1():
    disease = numpyro.sample("disease",dist.Bernoulli(0.4))
    treatment = numpyro.sample("treatment",dist.Bernoulli(0.7))
    list_variables=[disease,treatment]
    list_variables_no_treatment=[disease]
    
    bloodpressure = numpyro.sample("bloodpressure",dist.Normal(non_linear_fn(jnp.array(random.sample(list_variables_no_treatment,np.random.randint(0,1))),np.random.randint(0,1)*treatment),1))
    list_variables.append(bloodpressure)
    list_variables_no_treatment.append(bloodpressure)
                                   
    weight = numpyro.sample("weight", dist.Normal(non_linear_fn(jnp.array(random.sample(list_variables_no_treatment,np.random.randint(0,2))),np.random.randint(0,1)*treatment), 1))
    list_variables.append(weight)
    list_variables_no_treatment.append(weight)
    
    heartattack=numpyro.sample("heartattack",dist.Normal(non_linear_fn(jnp.array(random.sample(list_variables_no_treatment,np.random.randint(1,3))),np.random.randint(0,1)*treatment),1))
    list_variables.append(heartattack)
    list_variables_no_treatment.append(heartattack)
    
    variables=[]
    list_variables=[disease,treatment,bloodpressure,weight,heartattack]
    
    for i in range(5,k):
        vi=numpyro.sample('variable'+str(i),dist.Normal(non_linear_fn(jnp.array(random.sample(list_variables,np.random.randint(0,i))),np.random.randint(0,1)*treatment), 1))
        variables.append(vi)
        list_variables.append(vi)
    
    vi=numpyro.sample('variable'+str(k),dist.Normal(non_linear_fn(jnp.array(random.sample(list_variables,np.random.randint(0,k))),treatment), 1))
    variables.append(vi)
    list_variables.append(vi)
    return list_variables                       `

Now, I would like to generate "k" models, so I was thinking of some form of for-loop where i specify each of the models, but it didn't work. Any ideas?