Init_to_mean functionality

numpyro
2

What was your output? For me, your code worked as expected:

image
image2264×1460 444 KB