Hi all,
I am a physicist and I use deep learning on physical systems, where usually the physics is linear/simple when using complex values. That’s why I try to use complex values ANN, and I already use a custom set of functions/layers to implement complex layers:

  
      
      GitHub
  
  
    

wavefrontshaping/complexPyTorch 75

A high-level toolbox for using complex valued neural networks in PyTorch  - wavefrontshaping/complexPyTorch


  
  
    
    
  
  


So far, all my functions work taking two arguments, one tensor for the real part, one for the imaginary part, I am now wondering if it is the best way to go.
I recently discovered that PyTorch does have one type of complex layer, as it allows complex FFTs (which is awesome by the way):
https://pytorch.org/docs/master/torch.html?highlight=fft#torch.fft 43
This takes a tensor with an additional last dimension of size 2.
My question is, is it better to keep two tensors as arguments, or one with an additional dimension?
The advantage of the 2 arguments scheme is that, as most of the complex functions simply require to make independent operations on the imaginary and real part, it allows passing directly the two tensors to the builtin (real) PyTorch function. The simplest example being the C-relu:
def complex_relu(input_r,input_i):
    return relu(input_r), relu(input_i)

The problem is that the syntax is then not simple (especially for cost functions as they take two complex vectors, then 4 tensors).
If I want to use the one tensor argument, as in torch.fft, I then have to separate the real and imaginary part:
def complex_relu(complex_input):
    return torch.stack(relu(complex_input[...,0]), relu(complex_input[...,1]), dim = -1)

Is it really optimal? It seems that is would make copies of the sliced tensors, is it memory efficient? Is there a better way to do it?