Allow dot_product_attention to support normalization functions other than softmax

[CoderPat]

As discussed in #1291, I think adding an argument to dot_product_attention (and sub-functions of it) to pass a custom normalization function would allow greater flexibility to the library.

One use-case of this would be passing the entmax function (https://github.com/deep-spin/entmax-jax) to build sparse transformers. Another use case is passing the identity function to extract the unnormalized attention weights (and doing the normalization outside).

What do you think? (ping @avital due to the previous discussion)