(self)
| 158 | class WeightInitializerTest(TestCase): |
| 159 | # pylint: disable-next=no-self-use |
| 160 | def test_none_fan(self): |
| 161 | scale = 1.0 |
| 162 | init: WeightInitializer = ( |
| 163 | WeightInitializer.default_config() |
| 164 | .set(fan=None, scale=scale, distribution="uniform") |
| 165 | .instantiate() |
| 166 | ) |
| 167 | weight_shape = [100, 100] |
| 168 | weight = init.initialize( |
| 169 | "weight", prng_key=jax.random.PRNGKey(1), shape=weight_shape, dtype=jnp.float32 |
| 170 | ) |
| 171 | std_err = 1 / (12 * np.sqrt(np.prod(weight_shape))) * scale * 2 |
| 172 | self.assertBetween(np.mean(weight), 0.0 - 6 * std_err, 0.0 + 6 * std_err) |
| 173 | assert_allclose(jnp.min(weight), -scale, rtol=1e-4) |
| 174 | assert_allclose(jnp.max(weight), scale, rtol=1e-4) |
| 175 | |
| 176 | |
| 177 | class PerGroupInitializerTest(TestCase): |
nothing calls this directly
no test coverage detected