[Bug]: Jax, Jaxlib, Numpyro included as dependencies unconditionally

[BwL1289]

What happened?

I believe these should have been added as optional extras, not required dependencies.

See https://github.com/meta-pytorch/botorch/blob/main/botorch/models/fully_bayesian.py#L43-L51 as an example.

Separately, it's a bit confusing to me why a torch project would rely on Jax 🤷

BoTorch Version

0.18.0

Python Version

No response

Operating System

No response

(Optional) Describe any potential fixes you've considered to the issue outlined above.

No response

Pull Request

None

Code of Conduct

• I agree to follow BoTorch's Code of Conduct

[saitcakmak]

Thanks for flagging. I'll put up a fix.

Separately, it's a bit confusing to me why a torch project would rely on Jax 🤷

We wanted to avoid this but unfortunately Pyro is significantly slower than NumPyro, which relies on JAX. We've seen ~25x speed up in fully bayesian model fitting when switching to NumPyro, which was too much to ignore.