Add out_sharding argument to dot_product_attention

[samanklesaria]

Many operations like dot and conv have an out_sharding argument to reconcile incompatible arguments when using explicit sharding. Currently, dot_product attention lacks such an argument.

Ambiguity about the out_sharding location only occurs in binary operations. Dot product attention has three of them: the $QK^T$ multiplication, the addition of the bias term, and the multiplication by $V$. However, there will be substantial overlap in the output shardings of these intermediate values. Given an additional s_sharding argument, everything is identifiable.