Update references to JAX's GitHub repo

JAX has moved from https://github.com/google/jax to https://github.com/jax-ml/jax

PiperOrigin-RevId: 704187612
Change-Id: I341eb31510e29122eff2c3d0e5fb2f4a548360e6

README.md

@@ -39,7 +39,7 @@ Taking a Bayesian approach to MMM allows an advertiser to integrate prior inform
 - Report on both parameter and model uncertainty and propagate it to your budget optimisation.
 - Construct hierarchical models, with generally tighter credible intervals, using breakout dimensions such as geography.
 
-The LightweightMMM package (built using [Numpyro](https://github.com/pyro-ppl/numpyro) and [JAX](https://github.com/google/jax)) helps advertisers easily build Bayesian MMM models by providing the functionality to appropriately scale data, evaluate models, optimise budget allocations and plot common graphs used in the field.
+The LightweightMMM package (built using [Numpyro](https://github.com/pyro-ppl/numpyro) and [JAX](https://github.com/jax-ml/jax)) helps advertisers easily build Bayesian MMM models by providing the functionality to appropriately scale data, evaluate models, optimise budget allocations and plot common graphs used in the field.
 
 ## Theory
 

docs/index.rst

@@ -16,7 +16,7 @@ Installation
 We have kept JAX as part of the dependencies to install when installing
 LightweightMMM, however if you wish to install a different version of JAX or 
 jaxlib for specific CUDA/CuDNN versions see 
-https://github.com/google/jax#pip-installation for instructions on installing 
+https://github.com/jax-ml/jax?tab=readme-ov-file#instructions for instructions on installing 
 JAX. Otherwise our installation assumes a CPU setup.
 
 The recommended way of installing lightweight_mmm is through PyPi: