Add out_sharding argument to call methods for layers with jax calls that support it#5102

samanklesaria wants to merge 1 commit intogoogle:maingoogle/flax:mainfrom

samanklesaria:output_shardingsamanklesaria/flax:output_shardingCopy head branch name to clipboard