Method build_embedder

axlearn/common/layers_test.py:1256–1262

(dim, num_embeddings, rng, **kwargs)

class EmbedTest(parameterized.TestCase):
    @staticmethod
    def build_embedder(dim, num_embeddings, rng, **kwargs):
        cfg = Embedding.default_config().set(name="embed", dim=dim, num_embeddings=num_embeddings)
        if kwargs:
            cfg = cfg.set(**kwargs)
        emb = cfg.instantiate(parent=None)
        state = emb.initialize_parameters_recursively(rng)
        return (emb, state)

    @parameterized.parameters(itertools.product((5, 7), (2, 16), (10, 100), (True, False)))
    def test_embed_lookup(self, seq_len, dim, num_embeddings, is_training):