Trouble Using Custom c++ Operator with torch.compile

Alvinosaur: <p>I have a working setup using my custom c++ operator with torch.jit.script. I want to experiment with the newer torch.compile. I’m using pytorch==2.4.0</p>
<p>After seeing some errors with torch.compile, I’ve been <a href="https://docs.google.com/document/d/1_W62p8WJOQQUzPsJYa7s701JXt0qf2OfLub2sbkHOaU/edit#bookmark=id.2gbl3r1j2eg" rel="noopener nofollow ugc">following this guide</a> to add FakeTensor support for my c++ operator. I’m getting this error that I don’t see reported clearly anywhere:</p>
<pre><code class="lang-auto">RuntimeError: register_fake(...): the operator my_namespace::my_operator already 
has an implementation for this device type via a pre-existing registration to 
DispatchKey::CompositeImplicitAutograd.CompositeImplicitAutograd operators do not 
need an fake impl; instead, the operator will decompose into its constituents and those
 can have fake impls defined on them.
</code></pre>
<p>How should I change my fake operator?<br>
Here’s what my file structure looks like:</p>
<pre><code class="lang-auto">main.py
folder1
    folder2
        folder3
            file1.py
            file2.py
</code></pre>
<p>I will provide a simplified version of the files. Based on the error,</p>
<p>Here’s what file1.py looks like:</p>
<pre><code class="lang-auto">import torch
from . import file2
# torch.ops.import_module("file2")  # &lt;-- This import doesn't work for some reason
@torch.compile(backend="eager")
def helper(...) -&gt; torch.
    return torch.ops.my_namespace.my_operator(...)
</code></pre>
<p>Here’s what file2 looks like:</p>
<pre><code class="lang-auto">import torch
torch.ops.load_library("path_to_mycustom_op.so")

@torch.library.register_fake("my_namespace::my_operator")
def fake_operator(...):
    return torch.zeros(...)
</code></pre>
<p>Here’s what my custom c++ operator cpp file looks like:</p>
<pre><code class="lang-auto">// CUDA implementation for forward and backward
TORCH_LIBRARY_IMPL(my_namespace, CUDA, m) {
    m.impl("my_operator_forward", &amp;my_operator_forward);   // CUDA forward
    m.impl("my_operator_backward", &amp;my_operator_backward); // CUDA backward
}

// Don't define backward pass here since this is for non-autograd case, like inference
TORCH_LIBRARY(my_namespace, m) {
    m.def("my_operator_forward", &amp;my_operator_forward_no_autograd);
}

// Autograd implementation using the custom function
TORCH_LIBRARY_IMPL(my_namespace, Autograd, m) {
    m.impl("my_operator_forward", my_operator_autograd);  // Wraps Autograd
}
</code></pre>

Alvinosaur: <p>After trying the <a href="https://pytorch.org/tutorials/advanced/python_custom_ops.html" rel="noopener nofollow ugc">other method of registering a fake kernel</a>, torch compile seems to run successfully. I deleted file2.py entirely, and updated file1.py to be the following:</p>
<pre><code class="lang-auto">import torch
torch.ops.load_library("path_to_mycustom_op.so")

@torch.library.custom_op("fake::helper", mutates_args=())
def helper(...) -&gt; torch.Tensor:
    return torch.ops.my_namespace.my_operator(...)

@helper.register_fake
def _(...) -&gt; torch.Tensor:
    return torch.zeros(...)
</code></pre>