Update PyTorch to Flax porting documentation to include information on weight initialisation

Stumbled across google/jax#4862 discussion on parameter initialisation differences between PyTorch and Flax.

Would it be worth adding a note in the documentation that highlights these weight init differences? Users porting training scripts may expect the initialisation to be the same but find training behaviour is different as a result of different initialisation.

The issue discusses Linear and Conv layers mainly, but there are differences in other layers such as embedding layers.

Would be willing to contribute this myself if you can point me to / provide a quick guide on contributing to docs. Thanks :)