Add out_sharding arguments to linear layers where supported#5156

jackopenn wants to merge 2 commits intogoogle:maingoogle/flax:mainfrom

jackopenn:conv-out-shardingsjackopenn/flax:conv-out-shardingsCopy head branch name to clipboard