Method test_single_group

axlearn/common/param_init_test.py:179–212

(
        self,
    )

class PerGroupInitializerTest(TestCase):
    # pylint: disable-next=no-self-use
    def test_single_group(
        self,
    ):
        init_cfg = WeightInitializer.default_config()
        init = init_cfg.instantiate()
        per_group_init = (
            PerGroupInitializer.default_config()
            .set(
                initializer=init_cfg,
                num_groups=1,
            )
            .instantiate()
        )

        shape = [3, 3, 4, 6]

        prng_key = jax.random.PRNGKey(123)
        init_args = dict(
            prng_key=prng_key,
            shape=shape,
            dtype=jnp.float32,
            axes=FanAxes(in_axis=-2, out_axis=-1),
        )
        weight_standard_init = init.initialize(
            "weight",
            **init_args,
        )
        weight_per_group_init = per_group_init.initialize(
            "weight",
            **init_args,
        )

        # no difference when using num_groups=1
        assert_allclose(weight_standard_init, weight_per_group_init)

    # pylint: disable-next=no-self-use
    def test_constant_init(