Native Dot product attention

[AakashKumarNain]

With the latest release, we now have a native implementation of dot product attention provided using jax.nn.dot_product_attention(...). This is good except for one thing. When building nn in jax using libraries like Equinox, we are used to implement functionalities that work for a single example (not a batch_size of 1), and then we vmap over the model to make it work a batch. With the batch axis included in the attention implementation, I have two options:

1. Rewrite everything as it works for a batch
2. Expand qkv, and then squeeze the output of attention

I can do either but it seems to have broken my mental model for using vmap because in second case what does even batch mean anymore?

[Cjkkkk]

Looking into this and will make the API accept no batch input so it works better with vmap.

[ASEM000]

Some operations like general convolution or lstm accepts batch dimension by default. Do you think all batch-accepting operations should be changed too?

[AakashKumarNain]

Nowhere in my comment I am advocating for any change, that is purely up to the JAX team. All I am asking is to keep the mental model for vmap intact for every use case.

[ASEM000]

I agree with your motivation, I usually write a wrapper that uses expand-squeeze pattern around these functions. I was actually thinking about suggesting something like adding a keyword argument like batched=True to avoid breaking changes. WDYT?

[AakashKumarNain]

I usually write a wrapper that uses expand-squeeze pattern around these functions

Yeah I can do that except it's a wast of two ops followed by a copy

I was actually thinking about suggesting something like adding a keyword argument like batched=True to avoid breaking changes. WDYT?

That sounds good to me. The batched kwarg should default to False in this case, and the end-user should decide when to turn it on

[ASEM000]

Yeah I can do that except it's a waste of two ops followed by a copy

AFAIK If you are worried about the cost of copying, then under jit this pattern should not incur copying, I think you can double check with the core devs. For reshaping, it should not be an expensive op for this case.

[kaixih]

Can you give a shot of this change: #22830?

[AakashKumarNain]

Thanks. Now that it is merged, I will do some tests on my end