Making `torch` and `jax` entries in `extras_require` in `setup.py` for UB, so that we don't have to install all three frameworks if we only need one. Keeping TF as an always-install dependency because we rely on it for datasets, but certain tf-* packages could be moved once we start including more jax models.
Making `torch` and `jax` entries in `extras_require` in `setup.py` for UB, so that we don't have to install all three frameworks if we only need one. Keeping TF as an always-install dependency because we rely on it for datasets, but certain tf-* packages could be moved once we start including more jax models.

PiperOrigin-RevId: 380846679

google/uncertainty-baselines