From d13761513918734bf63a4cb6c2387be0a1d566f2 Mon Sep 17 00:00:00 2001 From: Jake Harmon Date: Wed, 4 Dec 2024 15:46:04 -0800 Subject: [PATCH] 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 --- edward2/jax/nn/heteroscedastic_lib.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/edward2/jax/nn/heteroscedastic_lib.py b/edward2/jax/nn/heteroscedastic_lib.py index 1370863d..55dff456 100644 --- a/edward2/jax/nn/heteroscedastic_lib.py +++ b/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',