(WIP?) vectorized log_likelihood function for NumPyro

[aporsch1]

Description

Checklist

• Follows official PR format
• New features are properly documented
• Code style correct (follows pylint and black guidelines)

---

📚 Documentation preview 📚: https://arviz--2390.org.readthedocs.build/en/2390/

[aporsch1]

Hey, I looked at the checks that failed, and they are failing because they can't even find test cases. I don't think that is related to the updated code at all? Let me know if I am missing something, though.

[OriolAbril]

@virajpandya could you try it out and see how timing compares to the ~80 mins from the latest release and setting log_likelihood=False?

You can install the arviz version of this PR with:

pip install "arviz @ git+https://github.com/aporsch1/arviz"

[OriolAbril]

Hey, I looked at the checks that failed, and they are failing because they can't even find test cases. I don't think that is related to the updated code at all? Let me know if I am missing something, though.

The pylint checks are failing. These are the specific errors:

************* Module arviz.data.io_numpyro
arviz/data/io_numpyro.py:195:64: C0303: Trailing whitespace (trailing-whitespace)
arviz/data/io_numpyro.py:196:0: C0301: Line too long (105/100) (line-too-long)
arviz/data/io_numpyro.py:195:34: E0602: Undefined variable 'jax' (undefined-variable)

For the jax import, note that it is not a dependency of ArviZ (nor it should be) so it needs to be imported at runtime from inside the method itself. This is already done in the __init__ method for example: https://github.com/arviz-devs/arviz/blob/main/arviz/data/io_numpyro.py#L67