Gradcheck fails for custom activation function

anirudh_puligandla: <p>I am implementing this activation function:</p>
<pre><code class="lang-auto">class EOActivation(Function):
    '''
    Electro-optic activations as described in - 
        {Williamson, Ian AD, et al. "Reprogrammable electro-optic nonlinear 
        activation functions for optical neural networks." IEEE Journal of 
        Selected Topics in Quantum Electronics 26.1 (2019): 1-12.}
    '''
    @staticmethod
    def forward(Z: Tensor,
                alpha: Tensor,
                g: Tensor,
                phi_b: Tensor) -&gt; Tensor:
        '''
        Forward-pass of EO nactivation function
        Z: tensor, Input tensor
        alpha: tensor, parameter 'alpha'
        g: tensor, parameter 'g'
        phi_b: tensor, parameter 'phi_b'
        '''
        return 1j * torch.sqrt(1 - alpha) * torch.exp(
            -1j*0.5*g*torch.conj(Z)*Z - 1j*0.5*phi_b) * torch.cos(
                0.5*g*torch.conj(Z)*Z + 0.5*phi_b) * Z
    
    @staticmethod
    def setup_context(ctx, inputs, output):
        '''
        ctx: Context object
        inputs: Inputs are the inputs to forward()
        output: Output tensor of forward()
        '''
        # Save parameters and output of forward for backward pass
        input, alpha, g, phi_b = inputs
        ctx.save_for_backward(input, alpha, g, phi_b)
    
    @staticmethod
    def backward(ctx, grad_Z: Tensor) -&gt; Tensor:
        '''
        ctx: context object
        grad_Z: backpropagated gradient signal from (l+1)th layer
        '''
        # get the parameters and output field computed during forward-pass
        Z, alpha, g, phi_b = ctx.saved_tensors
        zR, zI = Z.real, Z.imag
        # df_dRe - Gradient w.r.t. real part of the input
        df_dRe = torch.sqrt(1 - alpha) * torch.exp((-0.5*1j)*g*(zR - 1j*zI)*(zR + 1j*zI) - (
            0.5*1j)*phi_b) * (zR*g*(zI - 1j*zR) * torch.sin(0.5*(zR**2)*g + 0.5*(
                zI**2)*g + 0.5*phi_b) + ((zR**2)*g + 1j*zR*zI*g + 1j) * torch.cos(
                    0.5*(zR**2)*g + 0.5*(zI**2)*g + 0.5*phi_b))
        #df_dIm - Gradient w.r.t. imaginary part of the input
        df_dIm = torch.sqrt(1 - alpha) * torch.exp((-0.5*1j)*g*(zR - 1j*zI)*(zR + 1j*zI) - (
            0.5*1j) * phi_b) * (zI*g*(zI - 1j*zR) * torch.sin(0.5*(zR**2)*g + 0.5*(
                zI**2)*g + 0.5*phi_b) + (zR*zI*g + 1j*(zI**2)*g - 1) * torch.cos(
                    0.5*(zR**2)*g + 0.5*(zI**2)*g + 0.5*phi_b))
        # Return the gradient and 'None' for parameters in forward()
        return (grad_Z*df_dRe).real - 1j*(grad_Z*df_dIm).real, None, None, None
</code></pre>
<p>where alpha, g and phi_b are constant parameters initialized as tensors with “requires_grad=False”. The function is applied in a layer class whose forward() looks like this:</p>
<pre><code class="lang-auto">def forward(self, Z: Tensor) -&gt; Tensor:
        '''
        Z: Input tensor from (l-1)th layer
        Z_out: Output tensor after forward propagation (activation)
        '''
        Z = EOActivation.apply(Z, self.alpha, self.g, self.phi_b)       
        if self.photodetect:
            Z = Z.square()
        if self.bias is not None:
            Z = Z + self.bias.unsqueeze(0)

        return Z
</code></pre>
<p>let’s say “photodetect” and “bias” are always “False”.</p>
<p>My problem is that the activation does not work as expected. When used as part of a neural network, it fails to learn. I am compraing the equations with another <a href="https://github.com/Xoreus/neuroptica/blob/b0b4fad56af6002c3d852583be86393901c00dd3/neuroptica/nonlinearities.py#L153" rel="noopener nofollow ugc">implementation</a> and the calculations seem to be correct. My gradcheck output looks like this:<br>
<div class="lightbox-wrapper"><a class="lightbox" href="https://discuss.pytorch.org/uploads/default/original/3X/6/0/60b680d7b8316018ecfc2627105a87417a4ab7ec.png" data-download-href="https://discuss.pytorch.org/uploads/default/60b680d7b8316018ecfc2627105a87417a4ab7ec" title="JpQWNES2"><img src="https://discuss.pytorch.org/uploads/default/original/3X/6/0/60b680d7b8316018ecfc2627105a87417a4ab7ec.png" alt="JpQWNES2" data-base62-sha1="dNyQCfVYVsZgd9rfTzLynqdADh2" width="690" height="498" data-dominant-color="2F2E2F"><div class="meta"><svg class="fa d-icon d-icon-far-image svg-icon" aria-hidden="true"><use href="#far-image"></use></svg><span class="filename">JpQWNES2</span><span class="informations">1449×1047 85.7 KB</span><svg class="fa d-icon d-icon-discourse-expand svg-icon" aria-hidden="true"><use href="#discourse-expand"></use></svg></div></a></div><br>
Looks like there is a sign mismatch. Can anyone please tell me if I am missing any properties in my implementation? Or if I need to take any additional care for handling the gradients?</p>

soulitzer: <p>This might be a useful read if you’re dealing with derivatives involving complex numbers <a href="https://pytorch.org/docs/stable/notes/autograd.html#autograd-for-complex-numbers" class="inline-onebox" rel="noopener nofollow ugc">Autograd mechanics — PyTorch 2.5 documentation</a></p>

anirudh_puligandla: <p>Thank you very much for the resource. It is valuable. I understand that the computation of derivatives of funcitons of complex numbers is different from how it is done in the reference package that I was using. For now, from my initial understanding, I made a quick fix to my derivatives which works but I do not know why. I need to do that derivations myself following this guide that you shared. I will mark your comment as the solution so that the topic can be closed. Thank you again.</p>

anirudh_puligandla: <p>After going through pytorch’s gradient computations I see that the gradient that should be returned by this layer is not:</p>
<pre><code class="lang-auto">(grad_Z*df_dRe).real - 1j*(grad_Z*df_dIm).real
</code></pre>
<p>but,</p>
<pre><code class="lang-auto">torch.conj(grad_Z) * 0.5 * (df_dRe + 1j * df_dIm) + grad_Z * 0.5 * torch.conj(df_dRe - 1j * df_dIm)
</code></pre>
<p>However, as the function used in forward computation is derivable and as it can be written using only pytorch ops, a custom backward implementation is not really necessary in this case.</p>