Comment by sundalia

This is very interesting. A few questions:

- Why do you think similar approaches never landed on jax? My guess is this is not that useful for the current optimizations in fashion (transformers)

- How would you convince jax to incorporate this?