Add support for `out_sharding` in more layers

[jackopenn]

Currntly only out_sharding is supported in the linear layers. We can the argument to other layers such as Conv and Embed