Method test_repeat

axlearn/common/repeat_test.py:164–281

(self, dtype, remat_spec, drop_output, num_layers_total, unroll)

        unroll=(True, False, 1, 2),
    )
    def test_repeat(self, dtype, remat_spec, drop_output, num_layers_total, unroll):
        batch_size, num_layers = 14, 4
        cfg = _Ensemble.default_config().set(name="test", num_layers=num_layers_total, dtype=dtype)
        cfg.repeat_layer.set(
            remat_spec=remat_spec,
            drop_output=drop_output,
            unroll=unroll,
        )
        layer: _Ensemble = cfg.instantiate(parent=None)
        self.assertEqual(
            PartitionSpec(None),
            layer.create_parameter_specs_recursively()["repeat_layer"]["layer"]["inc"].mesh_axes,
        )
        layer_params = layer.initialize_parameters_recursively(prng_key=jax.random.PRNGKey(1))
        logging.info("layer params=%s", layer_params)

        input_forward_state = layer.init_forward_state(batch_size)
        if num_layers_total == num_layers:
            method = "forward"
            inputs = dict(
                carry=jnp.arange(batch_size, dtype=dtype),
                forward_state=input_forward_state,
            )
        else:
            method = "forward_first_n"
            input_forward_state = _get_first_n(num_layers, input_forward_state)
            inputs = dict(
                carry=jnp.arange(batch_size, dtype=dtype),
                forward_state=input_forward_state,
                n=num_layers,
            )

        (carry, output_forward_state), output_collection = F(
            layer,
            prng_key=jax.random.PRNGKey(2),
            state=layer_params,
            inputs=inputs,
            method=method,
            is_training=True,
            drop_output_collections=(),
        )
        logging.info("forward_state=%s", output_forward_state)
        logging.info("output_collection=%s", output_collection)
        assert_allclose(carry, jnp.arange(num_layers, num_layers + batch_size, dtype=dtype))
        self.assertEqual(shapes(input_forward_state), shapes(output_forward_state))
        assert_allclose(
            output_forward_state["repeat_layer"]["layer"],
            jnp.reshape(
                jnp.arange(batch_size)[None, :] + jnp.arange(num_layers, dtype=dtype)[:, None],
                (num_layers, batch_size),
            ),
        )
        # Check output collection.
        self.assertEqual(
            OutputCollection(
                state_updates={
                    "repeat_layer": {
                        # State update values are stacked across layers.