Allow for specifying axis for optax.projections

[JesseFarebro]

Currently, you can't specify the axis to take the projection along. Would the Optax authors be open to contributions allowing for specifying the axis of the projections? This would be much easier than having to vmap the projections.

[fabianp]

Thanks Jesse! That makes sense to me. Looping in @mblondel who is the original author of the projection module