I figured out the shape error. However I am now getting this warning and the sampling is very slow:
/usr/local/lib/python3.10/dist-packages/pytensor/tensor/slinalg.py:391: LinAlgWarning: Ill-conditioned matrix (rcond=1.5831e-19): result may not be accurate.
I tried running with numpyro on a GPU in Colab but it just ran through the sampling without accepting any steps. I am guessing there is an issue with the compatibility of custom distributions and numpyro?