Yep that's right! A very common use-case is for "document masking" (i.e. variable length sequences), and that requires recomputing the mask on every iteration (which isn't "free", but is on the order of microseconds to milliseconds and not seconds).

What does "there" mean in this case :)

I’ll count it!

Getting different attention masks working for AstroPT (a proto-foundation model for astronomy github.com/Smith42/astr...), so much nicer to do it with Flex Attention vs custom CUDA kernels -- thank you for releasing it to the world 🫡

Kinda interesting to me that the books I obsessively read as an elementary schooler are still some of the most popular series today.

I think torch-xla is definitely usable if you don’t want to train anything particularly weird or use unusual parallelism schemes. See this tweet from Saining Xie’s lab on evaluating torchxla vs. Jax for their use case: x.com/tongpetersb/...

The other nice parts about TPUs is that Google gives much more of them out for free compared to GPUs. Arguably this reflects how much people want to use them, but I think it's been a great boon for the academic labs willing to go through the effort.

! What were you using it for?

A lot of PyTorch is about dealing with this stuff nowadays!

Out of curiosity, what kind of shapes are you typically looking at?

Are they actually using FlexAttention here? I didn't see it in the repo