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: 702886844

Author: jakeharmon8

edward2/jax/nn/heteroscedastic_lib.py

@@ -87,7 +87,7 @@ def _get_cov_layer_kernel_init(self):
         )
       # Equivalent to the default kernel init `lecun_normal()` when we set the
       # scaling factor `scale_layer_kernel_init_factor` to 1.
-      # https://github.com/google/jax/blob/main/jax/_src/nn/initializers.py#L440
+      # https://github.com/jax-ml/jax/blob/main/jax/_src/nn/initializers.py#L440
       cov_layer_kernel_init = nn.initializers.variance_scaling(
           self.cov_layer_kernel_init_scale,
           'fan_in',