Method test_extend_step

axlearn/common/encoder_test.py:171–281

(self, prefill_states: bool, prefix_zero: bool)

        prefix_zero=[True, False],
    )
    def test_extend_step(self, prefill_states: bool, prefix_zero: bool):
        hidden_dim = 12
        vocab_size = 24
        num_heads = 4
        num_layers = 2
        source_length = 11

        encoder = CausalEncoder.default_config().set(
            dim=hidden_dim,
            vocab_size=vocab_size,
            dropout_rate=0,
            attention_mask=CausalAttentionLogitBiasLayer.default_config(),
            emb=bert_embedding_config(type_vocab_size=1, max_position_embeddings=source_length),
            transformer=bert_transformer_config(num_layers=num_layers, num_heads=num_heads),
            param_init=DefaultInitializer.default_config().set(
                init_by_param_name={
                    PARAM_REGEXP_WEIGHT: WeightInitializer.default_config().set(
                        fan=None, scale=0.02, distribution="normal"
                    )
                }
            ),
            pad_token_id=0,
        )
        set_layer_norm_eps_recursively(encoder, 1e-5)

        layer = encoder.set(name="layer_test").instantiate(parent=None)
        batch_size = 3

        # We ignore padding ids (0) for now to simplify the mask generation process.
        if prefix_zero:
            prefix = jnp.zeros([batch_size, 1], dtype=jnp.int32)
        else:
            prefix = jax.random.randint(
                jax.random.PRNGKey(123), [batch_size, 1], minval=1, maxval=vocab_size - 1
            )
        input_ids = jax.random.randint(
            jax.random.PRNGKey(123),
            [batch_size, source_length - 1],
            minval=1,
            maxval=vocab_size - 1,
        )
        input_ids = jnp.hstack([prefix, input_ids])

        params = layer.initialize_parameters_recursively(prng_key=jax.random.PRNGKey(123))

        ref_hidden_states, _ = F(
            layer,
            is_training=False,
            prng_key=jax.random.PRNGKey(123),
            state=params,
            inputs=dict(
                input_ids=input_ids,
                input_segment_ids=input_ids != 0,
                positions=jnp.arange(input_ids.shape[-1])[None, :],
            ),
        )
        ref_hidden_states = ref_hidden_states["hidden_states"]