Attention, dynamic shapes and JIT

[dfdx]

Probably, a trivial question, but I can't the right keywords to find an answer.

On one hand, jax.jit requires all arrays to have static shape.

On another hand, attention in transformers operates on sequences of difference length. Of course, it's possible to pad all sequences to the maximum length, but if, for example, my maximum is 4096 tokens and real sequence is 30 tokens, then it's a huge waste of computational resources.

How do people solve this dilemma?

[PhilipVinc]

Incremental padding at sizes that are either power of 2, or multiples of something else.

[doctor-phil]

@PhilipVinc could you explain how incremental padding could be done without changing the size of the array?

[jakevdp]

If you're curious how people implement attention in JAX, you could start by looking at implementations such as this one in flax: https://github.com/google/flax/blob/829e004c6d5646f24bd048023a4407e3b4d5aa98/flax/experimental/nnx/examples/07_transformer.py#L181

[doctor-phil]

@jakevdp Thank you, I'll take a look at this and it might be helpful.

My question is not about attention, necessarily. I am just wondering if there is a systematic way to solve problems that require dynamic array sizes without padding arrays to the maximum length and doing a lot of extra calculations.

[jakevdp]

In general, the strategy for handling dynamic array sizes in JAX is, depending on the situation, to either dispatch a new JIT-compiled operation for each size, or to pad all entries to a maximum size. In practice, the extra calculations done in the padding strategy may not be as important as you fear, especially if you're running on an accelerator like GPU or TPU.

[dfdx]

Thank you!