Fix Jax complex weight initialization

Created by: PhilipVinc

Fix the fact that jax networks with complex weights have zero imaginary part.

Adds a test