Use of jax.mask in seq2seq linen example

[iislucas]

Hello,

I was reading https://github.com/google/flax/tree/master/linen_examples/seq2seq/#L312 and noticed that it uses jax.mask, but when I looked that up in the jax API documentation I couldn't find it. I did find an undocumented implementation here: https://jax.readthedocs.io/en/latest/_modules/jax/api.html?highlight=mask#

Does anyone have a sense of the mask function is doing? And/or what role is has in this example?

Thanks!

[myagues]

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

There is not much information about jax.mask other than some comment. I guess it is still a WIP, and there are still some rough edges. The tests have some basic examples, but seq2seq is the only fully working example I know of.

Does anyone have a sense of the mask function is doing?

As I understand it, jax.mask is a way to circumvent static shape limitation in XLA, managing to ignore padded cells in the computation. It has similar function as TF and PyTorch masking, but with a different implementation.

And/or what role is has in this example?

Specifically for the seq2seq example, your input query and answer have been padded to a max_len amount using encode_onehot. When the Encoder scans over the query (or Decoder scans over the answer), you do not want to take into account the padded cells as it would waste computing resources. Wrapping the model function with jax.mask ensures that LSTMs will be unrolled up until the mask limit you provide.

flax/linen_examples/seq2seq/train.py

Lines 308 to 320 in 8092211

	IN_SHAPES = [{'query': '(n, _)', 'answer': '(m + 1, _)'}]
	OUT_ELEM = f'(m, {CTABLE.vocab_size})'
	OUT_SHAPE = (OUT_ELEM, OUT_ELEM)
	def apply_model(batch, in_masks, out_masks, params, key, teacher_force=True):
	@functools.partial(jax.mask, in_shapes=IN_SHAPES, out_shape=OUT_SHAPE)
	def model_fn(example):
	logits, predictions = model(teacher_force=teacher_force).apply(
	{'params': params},
	example['query'],
	example['answer'],
	rngs={'lstm': key})
	return logits, predictions
	return jax.vmap(model_fn)([batch], dict(n=in_masks, m=out_masks))

Considering query = [batch_size, query_max_len, vocab_size] and answer = [batch_size, answer_max_len, vocab_size], we take advantage of jax.vmap to do single element parallel computation, so each element of batch is processed with corresponding in_mask, out_mask value (this is why IN_SHAPES does not take into account batch size).

Hope this helps!

[iislucas]

Thanks! This is super helpful. Until there is better documentation in jax for this, might be nice this explanation in the seq2seq example.?