Hi again,
So I came up with some code to replace critical parts of NNs with complex functions so that it works on two N,C,H,W tensors for 2D or two N,C tensors for 1D in the form
xr,xi = complexLayer(xr,xi)
The code is here 9
It works nicely, and without time cost except for the complex batch norm. It uses the 2x2 covariance matrix between the real and imaginary parts, so there is calculation performed in python.
I have three issues with that complex batch norm code:


It is slow. I really do not want to go into cpp code, especially regarding point 3 below, so is there any easy and/or automatic way to optimize the code, like numba style? I would perform some quantitative test when I have some time.


There seems to be some memory leak, as the memory usage of my GPU increases over time till it crashes. When I use the “naive” complex batch norm approach, where is do standard batch norm independently on the real and imaginary part, I do not have memory issue, the ram usage is constant. Is there a need to del/free something inside the forward to prevent memory issues?


More a remark than an issue, I did not see any significant improvement compared to the naive approach. It is true also in the paper 1 I took it from but I work on a physical system where quantities are complex, so I hope it will improve things. Still, I do not give up so soon, I will continue testing.