Custom gradient for an input that is not a tensor

AlexZ: <p>Can I have a custom gradient for an input that is not a tensor?<br>
In other words, I want to get rid of the following error when I pass a function to a self-written <code>torch.autograd.Function</code></p>
<pre><code class="lang-auto">function ... returned a gradient different than None at position 1, but the corresponding forward input was not a Variable
</code></pre>
<h2>
<a name="my-use-case-1" class="anchor" href="#my-use-case-1"></a>My Use-case</h2>
<p>Let me explain my use-case: I have a variety of linear operators that look more or less like this:</p>
<pre><code class="lang-python">class LinearOperator(torch.nn.Module):
    def __init__(self, tensor):
        super().__init__()

        self.tensor = tensor
        
    def forward(self, x):
        """
        For a linear operator A and a tensor x this implements x -&gt; A(x)
        """
        return something( self.tensor, x )

    def rmatvec(self, x):
        """
        This implements the transpose of A applied to a tensor: x -&gt; A^T(x)
        """
        raise somethingTranspose( self.tensor, x )
</code></pre>
<p>The functions <code>something</code> and <code>somethingTranspose</code> are implemented using PyTorch so it should be possible to backprop through them.</p>
<p>While I can store the variable <code>tensor</code> with which I can compute the linear mapping, I cannot - due to storage space restrictions - store the entire transformation matrix that represents the linear operator <code>A</code>.</p>
<p>I also have a forward function that takes a linear operator A and a tensor x and produces some result f(A, x).<br>
While I can backpropagate through the function f, I want to implement the gradient differently (i.e. more efficiently) by hand in the backward function.</p>
<p>The code then looks like this:</p>
<pre><code class="lang-python">class MyFunctionF(torch.autograd.Function):

    @staticmethod
    def forward(ctx, A, x):
        ctx.A = A   # Linear Operator
        ctx.x = x   # PyTorch tensor

        result = computeF(A, x)
        # result is a PyTorch tensor

        return result

    @staticmethod
    def backward(ctx, grad_output):
        grad_A = something_A(ctx.A, ctx.x, grad_output)
        grad_x = something_x(ctx.A, ctx.x, grad_output)

        return grad_A, grad_x
</code></pre>
<p>I am using these classes in the following way:</p>
<pre><code class="lang-python"># Both tensor and x are PyTorch Tensors
tensor = ...
x = ...

linOp = LinearOperator(tensor)
res = MyFunctionF.apply(linOp, x)

linOp1 = OtherLinearOperator(tensor)
res1 = MyFunctionF.apply(linOp1, x)

linOp2 = NextLinearOperator(tensor)
res2 = MyFunctionF.apply(linOp2, x)
</code></pre>
<p>Here I have a multiple different classes of linear operators (<code>LinearOperator</code>, <code>OtherLinearOperator</code>, <code>NextLinearOperator</code>, …).</p>
<p>The forward pass works fine, but during the backward pass I am getting the error<br>
<code>function MyFunctionFBackward returned a gradient different than None at position 1, but the corresponding forward input was not a Variable</code></p>
<p><strong>I guess the problem is that <code>grad_A</code> in the <code>backward</code> function is a tensor, but the input <code>A</code> in <code>forward</code> is not a tensor. Is this correct?</strong></p>
<h2>
<a name="workaround-2" class="anchor" href="#workaround-2"></a>Workaround</h2>
<p>A workaround is to construct the linear operator in the forward function of <code>MyFunctionF</code>, but this means that</p>
<ol>
<li>I have to implement the function f for every type of linear operator.</li>
<li>I have to backpropagate through the forward method of the linear operator as well by hand, although the gradient of the Linear Operator forward method can be computed with backprop.</li>
</ol>
<pre><code class="lang-python">class MyFunctionF(torch.autograd.Function):

    @staticmethod
    def forward(ctx, tensor, x):
        A = LinearOperator(tensor)
        ctx.A = A
        ctx.x = x

        result = computeF(A, x)

        return result

    @staticmethod
    def backward(ctx, grad_output):
        grad_tensor = something_tensor(ctx.A, ctx.x, grad_output)
        grad_x = something_x(ctx.A, ctx.x, grad_output)

        return grad_tensor, grad_x
</code></pre>
<p><strong>Therefore my question:</strong><br>
<strong>Is there a different way to implement linear operators in conjunction with custom backward functions?</strong></p>

soulitzer: <p>Have you looked into: <a href="https://github.com/cornellius-gp/linear_operator" class="inline-onebox" rel="noopener nofollow ugc">GitHub - cornellius-gp/linear_operator: A LinearOperator implementation for PyTorch</a>? Demo notebook: <a href="https://github.com/Balandat/linear_operator/blob/demo_nb/examples/LinearOperator_demo.ipynb" class="inline-onebox" rel="noopener nofollow ugc">linear_operator/LinearOperator_demo.ipynb at demo_nb · Balandat/linear_operator · GitHub</a></p>

AlexZ: <p>Hey <img src="https://discuss.pytorch.org/images/emoji/apple/slight_smile.png?v=12" title=":slight_smile:" class="emoji" alt=":slight_smile:" loading="lazy" width="20" height="20"></p>
<p>sorry for the rather long delay. I finally had time to look at this issue again and dig a little deeper.</p>
<p>I had a look at <code>cornellius-gp/linear_operator</code> or to be more precise, I looked at <a href="https://docs.gpytorch.ai/en/stable/lazy.html" rel="noopener nofollow ugc">gpytorch.lazy</a> which is a more up-to-date version of the code by the same authors.<br>
Sadly this implementation does not solve my problems, as these lazy tensors are not an instance of <code>torch.Tensor</code>:</p>
<pre><code class="lang-python">import torch
import gpytorch

isinstance(
    gpytorch.lazy.DiagLazyTensor(torch.tensor((1.,2.,3.))),
    torch.Tensor
) # False
</code></pre>
<p>This is a problem when one is extending <code>torch.autograd.Function</code> which only works properly with a <code>torch.Tensor</code>. Let me illustrate this with a minimal example:</p>
<p>Create a custom <code>torch.autograd.Function</code>:</p>
<pre><code class="lang-python">class ForwardBackwardTest(torch.autograd.Function):
    @staticmethod
    def forward(ctx, a):
        print('forward')
        return a * 2

    @staticmethod
    def backward(ctx, grad_output):
        print('backward')
        return torch.rand((3,3)) # dummy gradient
</code></pre>
<p>For a <code>torch.Tensor</code>, the forward and backward passes work as expected:</p>
<pre><code class="lang-python">t = torch.rand((3,3), requires_grad=True)

res_t = fbt.apply(t).sum()
res_t.backward()
t.grad  # Everything working as expected
</code></pre>
<h2>
<a name="gpytorch-lazy-tensors-1" class="anchor" href="#gpytorch-lazy-tensors-1"></a>gpytorch lazy tensors</h2>
<p>For a gpytorch lazy tensor already the forward pass fails, because apparently the output of a custom forward method has to be a <code>torch.Tensor</code>.</p>
<pre><code class="lang-python">d = gpytorch.lazy.DiagLazyTensor(torch.tensor((1.,2.,3.)))
d.requires_grad_(True)

res_d = fbt.apply(d).sum()

# TypeError                                 Traceback (most recent call last)
# &lt;ipython-input-119-34bd74c90a99&gt; in &lt;module&gt;
#       2 d.requires_grad_(True)
#       3 
# ----&gt; 4 res_d = fbt.apply(d).sum()
# 
# TypeError: ForwardBackwardTestBackward.forward: expected Tensor or tuple of Tensor (got DiagLazyTensor) for return value 0
</code></pre>
<p>When we return a <code>torch.Tensor</code> from the <code>forward</code> method, the approach fails in the <code>backward</code> method.</p>
<p>So it seems that in order to implement things like linear operators and lazy tensors, that work with custom <code>forward</code> and <code>backward</code> methods, we have to make them instances of <code>torch.Tensor</code>.</p>
<h2>
<a name="subclasing-torchtensor-2" class="anchor" href="#subclasing-torchtensor-2"></a>Subclasing <code>torch.Tensor</code>
</h2>
<p>Let’s create  a <code>LinearOperator</code> class that inherits from <code>torch.Tensor</code>:</p>
<pre><code class="lang-python">class LinearOperator(torch.Tensor):
    @staticmethod 
    def __new__(cls, *args, **kwargs):
        shapeProvider = torch.rand((3,3))

        instance = torch.Tensor._make_subclass(cls, 
            shapeProvider, 
            True
        )
        return instance
      
    def __init__(self):
        pass

    def __mul__(self, t): 
        return torch.rand((3,3)) * t   # dummy implementation
</code></pre>
<p>In this implementation, the <code>.shape</code> property of the linear operator is set by PyTorch according to the <code>shape</code> of the <code>shapeProvider</code> tensor. This is a problem because then I have to allocate memory for the <code>shapeProvider</code> tensor. However, one of the main points of the linear operator approach is that I cannot store such a large tensor when working with very high-dimensional spaces. However, many linear operators between such spaces can easily be represented and stored.</p>
<p>Trying to reset the <code>.shape</code> property yields the following error:</p>
<pre><code class="lang-auto">AttributeError: attribute 'shape' of 'torch._C._TensorBase' objects is not writable
</code></pre>
<p>Using a <em>wrong</em> <code>.shape</code> for the linear operator is also not working, because PyTorch for example accesses the <code>shape</code> to check whether the gradient in the backward pass has the correct size.</p>
<h2>
<a name="conclusion-3" class="anchor" href="#conclusion-3"></a>Conclusion</h2>
<p>I think that there is currently no way to implement linear operator or lazy tensors in PyTorch, that</p>
<ul>
<li>do not need to store a tensor of the size of the operation they are representing</li>
<li>and work with a custom <code>torch.autograd.Function</code>.</li>
</ul>
<p>I would highly appreciate it if someone finds a way to do this.</p>