ONNX export fails for custom interpolate function with dynamic scale factor in torch.export.export (PyTorch 2.9.0)

sunxiaofan21: <p>Here’s a well-structured GitHub issue for the PyTorch repository based on your problem:</p>
<hr>
<h2><a name="p-476346-title-onnx-export-fails-for-custom-interpolate-function-with-dynamic-scale-factor-in-torchexportexport-pytorch-290-1" class="anchor" href="#p-476346-title-onnx-export-fails-for-custom-interpolate-function-with-dynamic-scale-factor-in-torchexportexport-pytorch-290-1"></a><strong>Title: ONNX export fails for custom interpolate function with dynamic scale factor in torch.export.export (PyTorch 2.9.0)</strong></h2>
<h3><a name="p-476346-description-2" class="anchor" href="#p-476346-description-2"></a><strong>Description</strong></h3>
<p>I’m encountering an issue when exporting a model to ONNX using <code>torch.onnx.export</code> with a custom interpolation function that takes a dynamic scale factor. The export fails during the <code>torch.export.export</code> step with a type error in <code>upsample_bicubic2d</code>.</p>
<h3><a name="p-476346-environment-3" class="anchor" href="#p-476346-environment-3"></a><strong>Environment</strong></h3>
<ul>
<li>PyTorch Version: 2.9.0+cu128</li>
<li>ONNX opset: 18</li>
</ul>
<h3><a name="p-476346-steps-to-reproduce-4" class="anchor" href="#p-476346-steps-to-reproduce-4"></a><strong>Steps to Reproduce</strong></h3>
<ol>
<li>Define a custom <code>torch.autograd.Function</code> for interpolation with a dynamic scale factor.</li>
<li>Use this function in a model’s <code>forward</code> method.</li>
<li>Attempt to export the model to ONNX using <code>torch.onnx.export</code>.</li>
</ol>
<p><strong>Minimal Reproduction Code:</strong></p>
<pre data-code-wrap="python"><code class="lang-python">import torch
from torch import nn
from torch.nn.functional import interpolate
import torch.onnx

class NewInterpolate(torch.autograd.Function):
    @staticmethod
    def symbolic(g, input, scales):
        return g.op(
            "Resize",
            input,
            g.op("Constant", value_t=torch.tensor([], dtype=torch.float32)),
            scales,
            coordinate_transformation_mode_s="pytorch_half_pixel",
            cubic_coeff_a_f=-0.75,
            mode_s="cubic",
            nearest_mode_s="floor"
        )

    @staticmethod
    def forward(ctx, input, scales):
        return interpolate(input, scale_factor=scales.tolist()[-2:], mode="bicubic", align_corners=False)

class StrangeSuperResolutionNet(nn.Module):
    def __init__(self):
        super().__init__()
        self.conv1 = nn.Conv2d(3, 64, kernel_size=9, padding=4)
        self.conv2 = nn.Conv2d(64, 32, kernel_size=1, padding=0)
        self.conv3 = nn.Conv2d(32, 3, kernel_size=5, padding=2)
        self.relu = nn.ReLU()

    def forward(self, x, upscale_factor):
        x = NewInterpolate.apply(x, upscale_factor)
        out = self.relu(self.conv1(x))
        out = self.relu(self.conv2(out))
        out = self.conv3(out)
        return out

model = StrangeSuperResolutionNet()
model.eval()
factor = torch.tensor([1, 1, 3, 3], dtype=torch.float32)
x = torch.randn(1, 3, 256, 256)

torch.onnx.export(
    model, 
    (x, factor), 
    "srcnn2.onnx", 
    opset_version=18, 
    input_names=['input', 'scale_factor'],
    output_names=['output']
)
</code></pre>
<h3><a name="p-476346-error-log-5" class="anchor" href="#p-476346-error-log-5"></a><strong>Error Log</strong></h3>
<pre><code class="lang-auto">[torch.onnx] Obtain model graph for `StrangeSuperResolutionNet` with `torch.export.export(..., strict=False)`... ❌
[torch.onnx] Obtain model graph for `StrangeSuperResolutionNet` with `torch.export.export(..., strict=True)`... ❌
Traceback (most recent call last):
  ...
TypeError: upsample_bicubic2d() received an invalid combination of arguments - got (FakeTensor, NoneType, bool, list), but expected one of:
 * (Tensor input, tuple of ints output_size, bool align_corners, tuple of floats scale_factors)
      didn't match because some of the arguments have invalid types: (FakeTensor, NoneType, bool, list of [SymInt, SymInt])
 * (Tensor input, tuple of ints output_size, bool align_corners, float scales_h = None, float scales_w = None, *, Tensor out = None)
</code></pre>
<h3><a name="p-476346-expected-behavior-6" class="anchor" href="#p-476346-expected-behavior-6"></a><strong>Expected Behavior</strong></h3>
<p>The model should export successfully to ONNX with the dynamic scale factor properly handled by the <code>torch.export.export</code> mechanism.</p>
<h3><a name="p-476346-additional-context-7" class="anchor" href="#p-476346-additional-context-7"></a><strong>Additional Context</strong></h3>
<ul>
<li>The issue appears to be related to type handling during the tracing phase of <code>torch.export.export</code>.</li>
<li>The same code might have worked in previous PyTorch versions with the older ONNX exporter.</li>
<li>The error occurs when <code>interpolate</code> is called with a list for <code>scale_factor</code> (from <code>scales.tolist()[-2:]</code>) while tracing with fake tensors.</li>
</ul>
<h3><a name="p-476346-question-8" class="anchor" href="#p-476346-question-8"></a><strong>Question</strong></h3>
<p>What is the recommended approach in PyTorch 2.9 for exporting models with dynamic resize/scale factors to ONNX, particularly when using the new <code>torch.export.export</code> infrastructure?</p>