Gradcheck fails for custom function

ssoni: <p>Hello all,</p>
<p>I’m new to PyTorch so this might turn out to be an easy fix, but I’m not able to get it working.</p>
<p>I’m trying to implement a custom <code>logsigmoid</code> function based on table lookup. I have the following code for my custom function:</p>
<pre><code class="lang-python">class LogSigmoid(Function):
    @staticmethod
    def forward (ctx, tensor):
        ctx.save_for_backward(tensor)
        return TableLookup.logsigmoid (tensor)

    @staticmethod
    def backward (ctx, grad_output):
        tensor, = ctx.saved_tensors
        grad_input = 1 - TableLookup.sigmoid(tensor)
        return grad_output * grad_input

from torch.autograd import gradcheck
input = (torch.randn(3,3,dtype=torch.double,requires_grad=True),)
test = gradcheck(LogSigmoid.apply, input, eps=1e-6, atol=1e-4)
print(test)
</code></pre>
<p>This fails with the error message:</p>
<pre><code class="lang-auto">RuntimeError: Jacobian mismatch for output 0 with respect to input 0,
numerical:tensor([[0., 0., 0., 0., 0., 0., 0., 0., 0.],
        [0., 0., 0., 0., 0., 0., 0., 0., 0.],
        [0., 0., 0., 0., 0., 0., 0., 0., 0.],
        [0., 0., 0., 0., 0., 0., 0., 0., 0.],
        [0., 0., 0., 0., 0., 0., 0., 0., 0.],
        [0., 0., 0., 0., 0., 0., 0., 0., 0.],
        [0., 0., 0., 0., 0., 0., 0., 0., 0.],
        [0., 0., 0., 0., 0., 0., 0., 0., 0.],
        [0., 0., 0., 0., 0., 0., 0., 0., 0.]], dtype=torch.float64)
analytical:tensor([[0.7356, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000],
        [0.0000, 0.4344, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000],
        [0.0000, 0.0000, 0.4974, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000],
        [0.0000, 0.0000, 0.0000, 0.2663, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000],
        [0.0000, 0.0000, 0.0000, 0.0000, 0.3699, 0.0000, 0.0000, 0.0000, 0.0000],
        [0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.6546, 0.0000, 0.0000, 0.0000],
        [0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.4759, 0.0000, 0.0000],
        [0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.2895, 0.0000],
        [0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.3435]],
       dtype=torch.float64)
</code></pre>
<p>If I replace the call to <code>TableLookup.logsigmoid</code> by <code>torch.nn.functional.logsigmoid</code>, then <code>gradcheck</code> works as expected. This led me to think that my lookup-based implementation might be incorrect, but I calculated the output using both <code>torch.nn.functional.logsigmoid</code> and <code>LogSigmoid.apply</code> and they are pretty much the same.</p>
<pre><code class="lang-python">from torch.autograd import gradcheck
input = (torch.randn(3,3,dtype=torch.double,requires_grad=True),)
print (F.logsigmoid(input[0]))
print (LogSigmoid.apply(input[0]))
test = gradcheck(LogSigmoid.apply, input, eps=1e-6, atol=1e-4)
</code></pre>
<pre><code class="lang-auto">tensor([[-0.2201, -1.0244, -0.1658],
        [-0.4461, -0.2707, -0.3701],
        [-0.9809, -0.5292, -0.3277]], dtype=torch.float64,
       grad_fn=&lt;LogSigmoidBackward&gt;)
tensor([[-0.2201, -1.0245, -0.1659],
        [-0.4461, -0.2707, -0.3701],
        [-0.9809, -0.5292, -0.3277]], dtype=torch.float64,
       grad_fn=&lt;LogSigmoidBackward&gt;)
</code></pre>
<p>For completeness, the <code>TableLookup</code> class is as follows:</p>
<pre><code class="lang-python">def sigm (x): return 1/(1+np.exp(-x))

class TableLookup (object):
    TABLE_SIZE = 100000
    MAX_EXP = 10
    
    EXP_TABLE = sigm ((np.arange (TABLE_SIZE) / TABLE_SIZE * 2 - 1) * MAX_EXP) #sigmoid table
    LOG_TABLE = np.log (EXP_TABLE) #logsigmoid table
    EXP_TABLE = torch.tensor (EXP_TABLE, dtype=torch.double) # TODO: how to do this is in the *device* memory?
    LOG_TABLE = torch.tensor (LOG_TABLE, dtype=torch.double)
    
    @classmethod
    def sigmoid (cls, tensor):
        """ For a given query tensor, calculate the sigmoid values"""
        # clamp
        tensor = torch.clamp (tensor, -TableLookup.MAX_EXP, ((TableLookup.TABLE_SIZE-1) / TableLookup.TABLE_SIZE * 2 - 1) * TableLookup.MAX_EXP)
        # map to indices
        indices = ((tensor + TableLookup.MAX_EXP) * TableLookup.TABLE_SIZE/2/TableLookup.MAX_EXP).long()
        # lookup
        return TableLookup.EXP_TABLE[indices]
        
    @classmethod
    def logsigmoid (cls, tensor):
        """ For a given query tensor, calculate the logsigmoid values as:
        logsigmoid (x) = x              if x &lt; -MAX_EXP
                       = lookup (x)     if x &gt;= -MAX_EXP and x &lt; MAX_EXP
                       = 0              if x &gt;= MAX_EXP
        """
        # if x &lt; -TableLookup.MAX_EXP:
        mask = tensor &lt; -TableLookup.MAX_EXP
        t1 = tensor * mask.double()
        
        tensor = torch.clamp (tensor, -TableLookup.MAX_EXP, ((TableLookup.TABLE_SIZE-1) / TableLookup.TABLE_SIZE * 2 - 1) * TableLookup.MAX_EXP)
        # map to indices
        indices = ((tensor + TableLookup.MAX_EXP) * TableLookup.TABLE_SIZE/2/TableLookup.MAX_EXP).long()
        # lookup
        t2 = TableLookup.LOG_TABLE[indices]
        # everywhere the mask is 1 should be converted to zero in t2
        t2[mask] = 0
        return t1 + t2.double()
</code></pre>
<p>Can someone explain what I might be doing wrong and how I may be able to fix it?</p>