Refactor to use complex numbers

[astanziola]

According to this repository for Fourier Neural Operators, they've made their code 30% faster by switching to native complex representation in PyTorch, which wasn't available when this work was published.

Is it worth refactoring the code to use complex multiplications? On one hand, it would be nice to have a faster codebase, on the other hand this repo makes sense for research but not really for production, so perhaps one doesn't really care about speed.

@sio13 @SonyPony Any suggestion? (Is it about time to just move everything to JAX? 😄 )

[SonyPony]

Well, FNO has way more complex multiplications than helmnet. However, even ~10% speedup would be nice. At least the training won't take a whole day 😛.

If the only part using complex operations is the residual computation, the refactoring should be quite straightforward.

JAX would be cool 😄, but I don't know whether it's worth moving the whole project to JAX.

[astanziola]

I think it is just that, probably only complex_mul needs to be removed and its call refactored.

Although I'm not totally sure that gradients are totally the same, see this note about autograd with complex numbers. Probably worth to add a quick test to be sure.