Exporting torch model with DFT results in transform being taken along the wrong dimension

[jamied157]

I'm using onnxscript==0.5.4 and torch==2.9.0 but expect this to still be a problem with the latest onnxscript as the code here: https://github.com/microsoft/onnxscript/blob/c1bfdfc5baf2dabf09191ad2ab7603ea21cfe614/onnxscript/function_libs/torch_lib/ops/fft.py hasn't been updated for a while.

Exporting a model in pytorch like

import torch.nn as nn
import torch

in_data = torch.randn(2, 128, 512)

class DFTModel(nn.Module):
    def __init__(self):
        super().__init__()
        
    def forward(self, a):
        return torch.fft.rfft(a)


torch.onnx.export(DFTModel(), (in_data,), "dft.onnx", dynamo=True)

Results in a model that looks like:

Whilst the output shapes are correct here, the attributes of the model are not. We unsqueeze the final axis and then set the dft_length to 2. This results in the DFT being calculated on the wrong dimension (the extra trailing dimension with padding). I think the correct graph should look more like:

To verify the model is wrong you can try and run it, I get this result

Whereas I get this with the correct model

Happy to go in and patch this with some pointers!

[justinchuby]

Could you verify if #2707 fixes the issue?

[justinchuby]

When I look at the implementation we don’t supply dft_length in fft_r2c. Not sure where it is coming from 🤔

[jamied157]

I tested with onnxscript @ git+https://github.com/microsoft/onnxscript.git@8a6176d6c46f9585d370bf3dcdbd5a1d28351d92 which is the commit corresponding to that branch. It did seem to work!

Weirdly the dft_length parameter has disappeared and the axis parameter has reverted to being positive (which I don't think should happen from the copilot MR). I did also try with an input size of 128x512 and that worked too

So maybe torch is doing something too? I did upgrade to 2.9.1

[jamied157]

When I look at the implementation we don’t supply dft_length in fft_r2c. Not sure where it is coming from 🤔

Is it possible onnxscript is accidentally setting the dft_length input instead of the axis input? That would explain things I think

[jamied157]

Oh it's this: 69025f7 🤦

[jamied157]

Will verify that this fixes but I expect so!

[jamied157]

Yep all sorted on onnxscript 0.5.6 thanks!