Error in mcmc.print_summary()

Is anyone familiar with the following error when calling mcmc.print_summary()? “TypeError: cannot perform reduce with flexible type”. I could find any information about it online.

The error is related to printing the divergences (in numpyro/infer/mcmc.py)

655         if "diverging" in extra_fields:
656             print(
--> 657                 "Number of divergences: {}".format(jnp.sum(extra_fields["diverging"]))
658             )

and the error is coming from jax numpy and eventually numpy itself.

I get this for any model I run. It’s not a huge problem because the summary table is printed nicely before the error occurs. It just doesn’t look very nice in my jupyter notebook.

This seems like a regression. Currently, our tests do not catch this issue. Could you make a colab notebook to illustrate it - and to double check that this is agnostic to the system that you run?