Add inputs_k and inputs_v args to attention layer

[chiamp]

Currently, MultiHeadDotProductAttention layer's call method signature is MultiHeadDotProductAttention.__call__(inputs_q, inputs_kv, mask=None, deterministic=None). As discussed in #1737, there are some cases where passing in separate values for the key and values is desired, which isn't possible with the current API. This PR adds two more arguments, inputs_k and inputs_v to the call method signature and sets the method signature to the following: MultiHeadDotProductAttention.__call__(inputs_q, inputs_k=None, inputs_v=None, *, inputs_kv=None, mask=None, deterministic=None). Note that the inputs_kv, mask and deterministic args are now keyword arguments.

• if inputs_k and inputs_v are None, then they will both copy the value of inputs_q (i.e. self attention)
• if inputs_v is None, it will copy the value of inputs_k (same behavior as the previous API, i.e. module.apply(inputs_q=query, inputs_k=key_value, ...) is equivalent to module.apply(inputs_q=query, inputs_kv=key_value, ...))
• if inputs_kv is not None, both inputs_k and inputs_v will copy the value of inputs_kv

Users can still use inputs_kv but a DeprecationWarning will be raised and inputs_kv will be removed in the future.
Since self attention can be done using this new API, the SelfAttention layer will also raise a DeprecationWarning and will be removed in the future.

Check out #3389 to see examples of how to port your code over to the new API.

[codecov-commenter]

Codecov Report

Merging #3379 (1d41190) into main (f20aed4) will increase coverage by 0.02%.
Report is 2 commits behind head on main.
The diff coverage is 90.90%.

@@            Coverage Diff             @@
##             main    #3379      +/-   ##
==========================================
+ Coverage   83.60%   83.62%   +0.02%     
==========================================
  Files          56       56              
  Lines        6746     6767      +21     
==========================================
+ Hits         5640     5659      +19     
- Misses       1106     1108       +2     

Files	Coverage Δ
flax/linen/attention.py	94.19% <90.90%> (-0.59%)	⬇️

... and 1 file with indirect coverage changes

[cgarciae]

Left a comment. Otherwise, looks good!

[chiamp]

Closing after this commit landed.