Copybara import of the project:
--
f6a222c710f87efbc3c1132af5deb853fb207ad9 by Marcus Chiam <marcuschiam@google.com>:
split inputs_kv arg in attention layer
COPYBARA_INTEGRATE_REVIEW=https://github.com/google/flax/pull/3379 from chiamp:attention f6a222c710f87efbc3c1132af5deb853fb207ad9
PiperOrigin-RevId: 572671273
