Use output of nn.Embed as a Layer

[pharringtonp19]

@matthias-wright provides a nice example of how to use nn.Embed -- here.

I would like to extend his example by using the output of nn.Embed as a linear layer --> Something along the lines of the following pseudo code.

class Net(nn.Module):
    
    @nn.compact
    def __call__(self, c, x):
        emb = nn.Embed(num_embeddings=10, features=5)(c)
        x = emb(x)
        return x

[marcvanzee]

In your example you have a 10x5 embedding table (i.e., 10 embeddings of size 5). emb contains the embeddings of your input c using that embedding table (so it is an array), so emb(x) doesn't mean anything.

Do you mean you want to add a linear layer after the embedding layer? Then you should use Dense and pass it emb.

[pharringtonp19]

@marcvanzee Thanks for the quick response.

I would like to use emb as the parameters/weights in the layer. Rather than passing emb to Dense, I was hoping to use emb as the parameters/weights of Dense and pass in x.

[marcvanzee]

In that case you have to replace the parameters of the Linear layer's kernel with emb. Something like this:

from flax import linen as nn
from jax import random
import jax.numpy as jnp
from flax.core.frozen_dict import unfreeze

x = jnp.zeros((2,), jnp.int32)

emb = nn.Embed(num_embeddings=10, features=5)
emb_variables = emb.init(random.PRNGKey(0), x)
emb_output = emb.apply(emb_variables, x)

dense = nn.Dense(features=5)
dense_variables = unfreeze(dense.init(random.PRNGKey(0), x))
dense_variables['params']['kernel'] = emb_output

dense_output = dense.apply(dense_variables, x)

If you want to replace the parameters inside the __call__ method, then you could consider using self.put_variable('params', 'kernel', value) to overwrite it inside the collection. However, be aware that you should mark the params collection as mutable during apply because by default we assume parameters are immutable, so you will get an error.

[marcvanzee]

Note the example above still includes a "bias" term in the Dense that isn't updated. If you don't want this you should pass use_bias=False to the Dense constructor.

[pharringtonp19]

@marcvanzee Thanks! This is great.