Pytorch "Extending Pytorch" doc missing ctx in forward method?

Sam-gege: <p>As stated, in the <a href="https://pytorch.org/docs/stable/notes/extending.html" rel="noopener nofollow ugc">Extending Pytorch</a> doc, the example:</p>
<pre><code class="lang-python">import torch
from torch.autograd import Function
# Inherit from Function
class LinearFunction(Function):

    # Note that forward, setup_context, and backward are @staticmethods
    @staticmethod
    def forward(input, weight, bias):
        output = input.mm(weight.t())
        if bias is not None:
            output += bias.unsqueeze(0).expand_as(output)
        return output

    @staticmethod
    # inputs is a Tuple of all of the inputs passed to forward.
    # output is the output of the forward().
    def setup_context(ctx, inputs, output):
        input, weight, bias = inputs
        ctx.save_for_backward(input, weight, bias)

    # This function has only a single output, so it gets only one gradient
    @staticmethod
    def backward(ctx, grad_output):
        ...
</code></pre>
<p>will cause an error:</p>
<pre><code class="lang-python">input=torch.randn(1,1)
weight=torch.randn(1,1,requires_grad=True)
bias=None
y=LinearFunction.apply(input,weight,bias)
# TypeError: forward() takes 3 positional arguments but 4 were given
</code></pre>
<p>Looks like you need to add <code>ctx</code> to the forward method:</p>
<pre><code class="lang-python">@staticmethod
    def forward(ctx, input, weight, bias):
        ...
</code></pre>
<p>it works afterwards</p>

Sam-gege: <p>even with the fix the backward is not working too, am I doing anything wrong?</p>
<pre><code class="lang-python">input=torch.randn(1,1)
weight=torch.randn(1,1,requires_grad=True)
bias=None
y=LinearFunction.apply(input,weight,bias)
y.backward()

Traceback (most recent call last):
  File "/home/.../Desktop/pytortto/src/test1.py", line 52, in &lt;module&gt;
    y.backward()
  File "/home/.../cpu_only/lib/python3.9/site-packages/torch/_tensor.py", line 488, in backward
    torch.autograd.backward(
  File "/home/.../cpu_only/lib/python3.9/site-packages/torch/autograd/__init__.py", line 197, in backward
    Variable._execution_engine.run_backward(  # Calls into the C++ engine to run the backward pass
  File "/home/.../cpu_only/lib/python3.9/site-packages/torch/autograd/function.py", line 267, in apply
    return user_fn(self, *args)
  File "/home/.../Desktop/pytortto/src/test1.py", line 30, in backward
    input, weight, bias = ctx.saved_tensors
ValueError: not enough values to unpack (expected 3, got 0)
</code></pre>

ptrblck: <p>The example works for me:</p>
<pre><code class="lang-python"># Option 1: alias
linear = LinearFunction.apply

x = torch.randn(10, 10)
lin = nn.Linear(10, 10)
out = linear(x, lin.weight, lin.bias)
out.mean().backward()
print(lin.weight.grad.abs().sum())
</code></pre>
<p>and you might want to check the <a href="https://pytorch.org/docs/stable/notes/extending.html#combined-or-separate-forward-and-setup-context">Combined or separate <code>forward()</code> and <code>setup_context()</code></a> section of the docs, which explains the usage of <code>setup_context</code> starting with PyTorch <code>2.0.0</code>.</p>

Sam-gege: <p>Yep I updated my pytorch from 1.13.1 to 2.0.0 and it works. Looks like <code>setup_context</code> only works starting from 2.0.0. Thanks!</p>