flax.linen.Embed documentation could be improved

[hawkinsp]

https://flax.readthedocs.io/en/latest/api_reference/flax.linen/_autosummary/flax.linen.Embed.html

• An example of using the module would be incredibly helpful. In fact, it would be incredibly helpful if every flax.linen module had an example in its documentation.
• The docs say:

A parameterized function from integers [0, n) to d-dimensional vectors.

num_embeddings
number of embeddings.

It wasn't immediately obvious to me that "number of embeddings" really meant "n", which is often called the vocabulary size.

• In

features
number of feature dimensions for each embedding.

it is also slightly unclear that this is the same thing as d. Either use the same symbol, or refer back to the name used above?

• I'll also note that the attend and setup methods are either undocumented or the links to their documentation aren't working.

[chiamp]

Sorry for the delay, #3453 added example usages in the docstrings of each Module layer.
#3539 will additionally add more detail to the Embed documentation.

[jmerskine1]

The link for the updated documentation seems broken. I could be wrong, but there doesn't seem to be any example either. Has this correction been redacted?

[chiamp]

Apologies, this is the new link: https://flax.readthedocs.io/en/latest/api_reference/flax.linen/layers.html#flax.linen.Embed

[jmerskine1]

@chiamp Thanks very much for your super quick response!