Method test_clip_textual_encoder

axlearn/vision/clip_test.py:111–154

(self, act_fn)

    @parameterized.parameters(["nn.gelu", "quick_gelu"])
    def test_clip_textual_encoder(self, act_fn):
        act_fn_key = act_fn.replace(".", "_")
        golden = _load_golden_jax(f"test_clip_textual_encoder_{act_fn_key}")

        model_dim = 32
        ff_dim = 64
        num_layers = 3
        num_heads = 8
        max_seq_len = 12

        kwargs = {
            "pad_token_id": OUR_PAD_TOKEN_ID,
            "max_seq_len": max_seq_len,
            "vocab_size": OUR_VOCAB_SIZE,
            "num_layers": num_layers,
            "model_dim": model_dim,
            "num_heads": num_heads,
            "feed_forward_dim": ff_dim,
            "feed_forward_act": act_fn,
            "dropout_rate": 0,
            "projection_dim": model_dim,
        }

        layer_cfg = set_text_encoder_config(**kwargs)
        layer_cfg.set(name="test")
        layer = layer_cfg.instantiate(parent=None)

        # Golden params cover the text_encoder subtree; initialize full params and overlay.
        layer_params = layer.initialize_parameters_recursively(prng_key=jax.random.PRNGKey(0))
        layer_params.update(golden["params"])
        # Convert HF-format input_ids (PAD==EOS==49407) to AXLearn format (PAD=49408).
        input_ids = _hf_to_axlearn_input_ids(golden["inputs"]["input_ids"])

        layer_outputs, _ = F(
            layer,
            inputs=dict(input_batch={"text": jnp.expand_dims(jnp.asarray(input_ids), 1)}),
            state=layer_params,
            is_training=True,
            prng_key=jax.random.PRNGKey(0),
        )
        assert_allclose(
            jnp.squeeze(layer_outputs[TEXT_EMBEDDINGS]),
            jnp.asarray(golden["outputs"]["pooler_output"]),
        )


class TestCLIPModel(parameterized.TestCase):