Using jax.mask with BiRNN for text classification

[myagues]

I have been trying to make jax.mask work with a Bidirectional RNN, but I am not sure where the mask decorator should be placed.
My code is a replication of the TF tutorial Text classification with an RNN, but using Flax, and the code can be run in Colab. The main model and its call are:

class TextModel(nn.Module):
    vocab_size: int
    emb_dim: int = 64
    hidden_size: int = 64

    @nn.compact
    def __call__(self, inputs):
        x = inputs.astype("int32")
        x = nn.Embed(num_embeddings=self.vocab_size, features=self.emb_dim)(x)

        bi_lstm = BiLSTM(self.hidden_size)
        x = bi_lstm(x)

        x = nn.Dense(64)(x)
        x = nn.relu(x)
        logits = nn.Dense(1)(x)
        return logits


def apply_model(inputs, params, model):
    def model_fn(example):
        logits = model.apply({"params": params}, example)
        return logits

    return jax.vmap(model_fn)(inputs)

and when trying to use a mask (as in the seq2seq example):

def apply_model(inputs, params, model):
    @functools.partial(jax.mask, in_shapes=["n"], out_shape="_")
    def model_fn(example):
        logits = model.apply({"params": params}, example)
        return logits

    masks = jnp.sum(jnp.not_equal(inputs, 0.0), axis=-1)
    return jax.vmap(model_fn)([inputs], dict(n=masks))

it returns an error with UndefinedPoly: Polynomial n is not constant that traces back to nn.Embed (guess jax.mask does not like the shape change). Are the shapes in the mask not correctly defined, or should I just define the mask around the BiLSTM module?

Thanks!!

[avital]

Hi @myagues -- jax.mask isn't used very broadly so it may still have some rough edges... @marcvanzee did the most work with masking so far. Marc, can you take a look at this?

[myagues]

It seems jax.mask is being redesigned and its use is currently not recommended as per this comment.

[avital]

Sadly, I wonder if this means we should step back from using mask in any of our examples... /cc @marcvanzee @mattjj

[mattjj]

Yes, but I wouldn't be sad about it! We're redesigning it to be better, after all.

[marcvanzee]

The seq2seq example is currently the only example in which we are using masking, we left it out of other examples because the masking API is still in flux. We did implement it in the seq2seq example because the JAX team was looking for a use-case to test out their early masking API, so at that time it was a win-win situation where we managed to catch a number of bugs.

If the current API is going away then I guess it doesn't make sense to keep using it in the seq2seq example. @mattjj, do you have an idea of when you think the current API is no longer supported, and our example will break?

[mattjj]

I'm not sure, but we would want to coordinate that timing with you so that everything keeps working smoothly on your end. Also, we don't plan to remove the API until we have a better replacement, and because we want to ensure the new API covers the previous use cases well I plan to work with you to update the seq2seq example. Your seq2seq example has already been really helpful for both the old and new APIs!