How does Haiku keep track of layers?

[vedantroy]

Taking a look at the transformer example: https://github.com/deepmind/dm-haiku/blob/main/examples/transformer/train.py, it seems like the layers are created using function calls. For example:

token_embedding_map = hk.Embed(vocab_size, d_model, w_init=embed_init)

How does Haiku maintain a consistent reference to this state over time? And let's say I want to, for example, extract out the embedding & save it to a file. How would I go about doing that? It's a bit hard to tell since there's no consistent object that acts as a "handle" to the Embed layer.