flax
Added out_sharding and qkv_sharding to MHA call method
#5434
Merged

Loading