Sparse token prediction with jit

[mk-0]

I'm new to JAX but I wasn't able to find an answer to this anywhere. I am building a BERT-like language model with MLM loss on about 15% of all input tokens.

Given hidden outputs h of shape [sequence, features], the output emedding W of shape [features, vocab], target y of shape [sequence], and dynamic boolean mask mask of shape [sequence] where about 15% of all values are True, I only need loss values for masked tokens, i. e.:
cross_entropy(softmax(h @ W), one_hot(y))[mask].mean().
Without JIT, I would save memory by doing h[mask] @ W instead, but with JIT, arrays of dynamic size are not allowed.

It seems important though, because the matrix [sequence, vocab] if of course quite large. Is there a way for me to make sure the optimizer does not produce unnecessary rows without creating arrays of dynamic shapes, given that mask is dynamic and has slightly different number of True values at each step? Or is the only way to change my strategy and mask exactly 15% of each sample instead, so that the shapes are static?