Add out_sharding arguments to linear layers where supported

[jackopenn]

What does this PR do?

• Add out_sharding arguments to linear layers where supported, these include:
  • Conv : jax.lax.conv_general_dilated
  • Embed.attend : jnp.dot
• Update tests for all linear layers and new methods

Fixes #5155

Checklist

• This PR fixes a minor issue (e.g.: typo or small bug) or improves the docs (you can dismiss the other checks if that's the case).
• This change is discussed in a Github issue/discussion (please add a link).
• The documentation and docstrings adhere to the documentation guidelines.
• This change includes necessary high-coverage tests. (No quality testing = no merge!)

[samanklesaria]

Thanks for the PR! This is stuff I should have added in #5102 - apologies for missing these modules. This all looks good to me!

[samanklesaria]

@jackopenn looks like there's some issues with the tests.

[jackopenn]

@jackopenn looks like there's some issues with the tests.

Oops. Fixed now :)