Sorry, I didn't check that there was already an open PR for this task #1727. I added a commit for bumping the jaxns version to the latest, and will close my PR if required

Thanks, @renecotyfanboy! Could you run make format and make lint to fix the lint issue?

I manually linted the file with ruff, the make lint and make format tools fails locally due to the header script :

python scripts/update_headers.py
Traceback (most recent call last):
  File "/Users/sdupourque/PycharmProjects/numpyro/scripts/update_headers.py", line 45, in <module>
    while lines[lineno].startswith(comment.format("Copyright")):
IndexError: list index out of range

Seems like there is a python file which has conflict format in your folder. CI works btw. :)