Almost got it to run through Bambi. But getting an XLA runtime error on compilation. I have no idea what that error means. I used the following at the start of the notebook:
import numpyro numpyro.set_platform(platform=‘METAL’)