MCPcopy Create free account
hub / github.com/apple/axlearn / test_scale

Method test_scale

axlearn/common/param_init_test.py:42–48  ·  view source on GitHub ↗
(self)

Source from the content-addressed store, hash-verified

40 np.testing.assert_array_equal(bias, jnp.zeros_like(bias))
41
42 def test_scale(self):
43 init: DefaultInitializer = DefaultInitializer.default_config().instantiate()
44 scale = init.initialize(
45 "scale", prng_key=jax.random.PRNGKey(1), shape=[4], dtype=jnp.bfloat16
46 )
47 self.assertEqual(scale.dtype, jnp.bfloat16)
48 np.testing.assert_array_equal(scale, jnp.ones_like(scale))
49
50 def test_init_by_param_name_add(self):
51 cfg = DefaultInitializer.default_config()

Callers

nothing calls this directly

Calls 3

instantiateMethod · 0.45
default_configMethod · 0.45
initializeMethod · 0.45

Tested by

no test coverage detected