Does my model make sense.?

Hi,
I have cooked my crude code in Colab HERE

I deals why the following simple exercise:

A) Make a galaxy image generator
1. generate an image of a galaxy profile 128x128 pix
2. convolve with a PSF
3. add a Poisson noise

Then generate an observation

B) Make a Numpyro Model with 2 parameters entering in the galaxy image prod (flux and sigma)

C) setup a NUTS sampling.

The points I am questionning are

  1. the Numpyro model ie. the likelood is it correct?
  2. the running of the NUTS is very slow (I’m on a Colab CPU): does it is due to my poor Numpyro modeling, or their is a trick to speed up (of course on GPU I get higher speed), or the problem is hard ?

I need some expertise. Thanks

Concerning point 2: the GPU helps and I was squeezing one of the parameters by the prior.
So the main question is the 1st one.