Error in official document code?

Aakira: <p>I’m following the tutorial <a href="https://pytorch.org/docs/stable/notes/extending.html#:~:text=class%20MyCube(torch.autograd.Function)%3A" rel="noopener nofollow ugc">here</a> and use the code:</p>
<pre><code class="lang-auto">class MyCube(torch.autograd.Function):
    @staticmethod
    def forward(x):
        # We wish to save dx for backward. In order to do so, it must
        # be returned as an output.
        dx = 3 * x ** 2
        result = x ** 3
        return result, dx

    @staticmethod
    def setup_context(ctx, inputs, output):
        x, = inputs
        result, dx = output
        ctx.save_for_backward(x, dx)

    @staticmethod
    def backward(ctx, grad_output, grad_dx):
        x, dx = ctx.saved_tensors
        # In order for the autograd.Function to work with higher-order
        # gradients, we must add the gradient contribution of `dx`,
        # which is grad_dx * 6 * x.
        result = grad_output * dx + grad_dx * 6 * x
        return result

# Wrap MyCube in a function so that it is clearer what the output is
def my_cube(x):
    result, dx = MyCube.apply(x)
    return result

x = torch.tensor(2.0, requires_grad=True)
y, dx = my_cube(x)  # y = x^3, dx = 3x^2
</code></pre>
<p>But I got this error:</p>
<pre><code class="lang-auto">File "H:\code\lab\lab.py", line 33, in &lt;module&gt;
    y, dx = my_cube(x)  # y = x^3, dx = 3x^2
  File "H:\code\lab\lab.py", line 29, in my_cube
    result, dx = MyCube.apply(x)
TypeError: MyCube.forward() takes 1 positional argument but 2 were given
</code></pre>

KFrank: <p>Hi Aakira!</p>
<aside class="quote no-group" data-username="Aakira" data-post="1" data-topic="216146" data-full="true">
<div class="title">
<div class="quote-controls"></div>
<img loading="lazy" alt="" width="24" height="24" src="https://discuss.pytorch.org/user_avatar/discuss.pytorch.org/aakira/48/58984_2.png" class="avatar"> Aakira:</div>
<blockquote>
<p>But I got this error:</p>
<pre><code class="lang-auto">File "H:\code\lab\lab.py", line 33, in &lt;module&gt;
    y, dx = my_cube(x)  # y = x^3, dx = 3x^2
  File "H:\code\lab\lab.py", line 29, in my_cube
    result, dx = MyCube.apply(x)
TypeError: MyCube.forward() takes 1 positional argument but 2 were given
</code></pre>
</blockquote>
</aside>
<p>This works* for me (using the latest stable version, 2.6.0).</p>
<p>Here is (a tweaked version of) your example script:</p>
<pre data-code-wrap="python"><code class="lang-python">import platform
print (platform.python_version())

import torch
print (torch.__version__)

class MyCube(torch.autograd.Function):
    @staticmethod
    def forward(x):
        # We wish to save dx for backward. In order to do so, it must
        # be returned as an output.
        dx = 3 * x ** 2
        result = x ** 3
        return result, dx

    @staticmethod
    def setup_context(ctx, inputs, output):
        x, = inputs
        result, dx = output
        ctx.save_for_backward(x, dx)

    @staticmethod
    def backward(ctx, grad_output, grad_dx):
        x, dx = ctx.saved_tensors
        # In order for the autograd.Function to work with higher-order
        # gradients, we must add the gradient contribution of `dx`,
        # which is grad_dx * 6 * x.
        result = grad_output * dx + grad_dx * 6 * x
        return result

# Wrap MyCube in a function so that it is clearer what the output is
def my_cube(x):
    result, dx = MyCube.apply(x)
    # return result
    return result, dx   # return a tuple containing both result and dx

x = torch.tensor(2.0, requires_grad=True)
y, dx = my_cube(x)  # y = x^3, dx = 3x^2   # this line asks for a tuple, so we modified my_cube()

print (x, y, dx)
</code></pre>
<p>And here is its output:</p>
<pre><code class="lang-plaintext">3.12.8
2.6.0+cu126
tensor(2., requires_grad=True) tensor(8., grad_fn=&lt;MyCubeBackward&gt;) tensor(12., grad_fn=&lt;MyCubeBackward&gt;)
</code></pre>
<p>Double-check for possible copy-paste errors and maybe check your pytorch (and<br>
python) version, but this should work with any relatively recent version of pytorch.</p>
<p>*)  The last line of your example script requires <code>my_cube()</code> to return a sequence<br>
of some sort (e.g., a tuple), so I tweaked <code>my_cube()</code> accordingly.</p>
<p>Best.</p>
<p>K. Frank</p>

Aakira: <p><a class="mention" href="/u/kfrank">@KFrank</a> Thank for taking your time.</p>
<p>My code just has not worked yet:</p>
<pre><code class="lang-auto">3.10.9
1.13.1+cu117
Traceback (most recent call last):
  File "H:\code\lab\lab.py", line 38, in &lt;module&gt;
    y, dx = my_cube(x)  # y = x^3, dx = 3x^2   # this line asks for a tuple, so we modified my_cube()
  File "H:\code\lab\lab.py", line 33, in my_cube
    result, dx = MyCube.apply(x)
TypeError: MyCube.forward() takes 1 positional argument but 2 were given
</code></pre>

KFrank: <p>Hi Aakira!</p>
<aside class="quote no-group" data-username="Aakira" data-post="3" data-topic="216146" data-full="true">
<div class="title">
<div class="quote-controls"></div>
<img loading="lazy" alt="" width="24" height="24" src="https://discuss.pytorch.org/user_avatar/discuss.pytorch.org/aakira/48/58984_2.png" class="avatar"> Aakira:</div>
<blockquote>
<pre><code class="lang-auto">3.10.9
1.13.1+cu117
Traceback (most recent call last):
  File "H:\code\lab\lab.py", line 38, in &lt;module&gt;
    y, dx = my_cube(x)  # y = x^3, dx = 3x^2   # this line asks for a tuple, so we modified my_cube()
  File "H:\code\lab\lab.py", line 33, in my_cube
    result, dx = MyCube.apply(x)
TypeError: MyCube.forward() takes 1 positional argument but 2 were given
</code></pre>
</blockquote>
</aside>
<p>This error makes it seem that <code>MyCube.forward()</code> is being treated as a non-static<br>
class method (where when you call <code>.forward (x)</code> with a single explicit argument,<br>
it actually gets called with an additional <code>self</code> argument: <code>.forward (self, x)</code>).</p>
<p>Quite some time ago, the <code>.forward()</code> method of <code>torch.autograd.Function</code> was<br>
non-static.  But you seem to be using pytorch 1.13, which uses static methods<br>
according to <a href="https://pytorch.org/docs/1.13/autograd.html#torch.autograd.Function" rel="noopener nofollow ugc">its documentation.</a></p>
<p>Could you post the entirety of your lab.py as well as the output you when you run it<br>
on the latest stable pytorch (2.6.0)?</p>
<p>Best.</p>
<p>K. Frank</p>