Repository: microsoft/onnxscript
Pull request #2141: Implement fft torchop
Author: login: bmehta001
Labels: name: module: torchlib
description: Related to the torch/aten function lib in development
Created: 2025-03-27T16:58:21Z | Merged: 2025-04-16T01:34:50Z

WIP

- Implement aten__fft_r2c, aten__fft_c2r, aten__fft_c2c
r2c = forwards, could be one-sided
c2r = backwards/inverse, never one-sided
c2c could be either forwards/backwards, never one-sided

Must respect normalization method provided - however, op.DFT calls "backwards" normalization, if 'inverse' is set to True, so need to account for normalization being done by op.DFT

When running above functions across multiple axes, need to run FFT in reverse order through op.DFT one-by-one

Currently have issues with:
- c2r has extra parameter of last_dim_size, so must truncate/zero-pad to ensure last dimension size matches last_dim_size -- still debugging this part to avoid triggering errors

https://github.com/microsoft/onnxscript/issues/1271



Diff:
diff --git a/onnxscript/function_libs/torch_lib/ops/fft.py b/onnxscript/function_libs/torch_lib/ops/fft.py
index 51621ed596..ea92dc347d 100644
--- a/onnxscript/function_libs/torch_lib/ops/fft.py
+++ b/onnxscript/function_libs/torch_lib/ops/fft.py
@@ -21,98 +21,33 @@
 from onnxscript.onnx_types import TensorType
 
 
-@torch_op(
-    ("aten::_fft_c2c", "aten::_fft_c2r", "aten::_fft_r2c"),
-    private=True,
-    complex=True,
-    trace_only=True,
-)
 def _fftn_onnx_normalization(
-    self,
-    transformed: TFloat,
+    self: TFloat,
     normalization: int,
-    forward: bool,
-    dims: Sequence[int],
-) -> TFloat:
-    # Obtain the total_sample_count (n) for normalization
-    self_shape = op.Shape(self)
-    total_sample_count = op.ReduceProd(op.Gather(self_shape, dims), keepdims=0)
-    total_sample_count = op.CastLike(total_sample_count, transformed)
-
-    # Normalize the result
-    # Reference https://pytorch.org/docs/stable/generated/torch.fft.fftn.html#torch.fft.fftn
-    # Reference https://github.com/pytorch/pytorch/blob/d090c18fcaaba6e1b5cb474a89058cf6081c8275/torch/_refs/fft.py#L42
-    if normalization == 1:
-        # "forward" - normalize by 1/n
-        if forward:
-            result = op.Div(transformed, op.Sqrt(total_sample_count))
-        else:
-            result = op.Mul(transformed, op.Sqrt(total_sample_count))
-    elif normalization == 2:
-        # "ortho" - normalize by 1/sqrt(n)
-        if forward:
-            result = op.Div(transformed, total_sample_count)
-        else:
-            result = transformed
-    else:
-        # "backward" - no normalization
-        if forward:
-            result = transformed
-        else:
-            result = op.Mul(transformed, total_sample_count)
-
-    return result
-
-
-@torch_op(
-    ("aten::_fft_c2c", "aten::_fft_c2r", "aten::_fft_r2c"),
-    trace_only=True,
-    private=True,
-    complex=True,
-)
-def _fftn_onnx(
-    self: TFloat, dims: Sequence[int], normalization: int, inverse: bool, onesided: bool
+    signal_size: INT64,
+    inverse: bool = False,
 ) -> TFloat:
-    """Standard complex to complex or real to complex FFT (forward or backward).
-
-    This is a private shared function for implementing the various FFT functions.
-
-    Args:
-        self: The input tensor.
-        dims: The dimensions to apply FFT.
-        normalization: The normalization mode.
-        inverse: Whether to compute the inverse FFT.
-        onesided: Whether to compute the one-sided FFT, which retains only the
-            positive frequencies.
-
-    Returns:
-        The transformed tensor.
-    """
-
-    # NOTE: trace_only because we need to process each dimension in a loop
-    # NOTE: SymInt dim is not support because DFT-17 needs a static axis
-    # TODO(justinchuby): Make dim dynamic and remove trace_only when ONNX provides support
-
-    # The 0-th dimension in ONNX DFT-17 is the batch dimension. We need to add a new
-    # dimension at the beginning to represent the batch dimension.
-    transformed = op.Unsqueeze(self, axes=[0])
-
-    # Add 1 to account for the batch dimension when counting axes from the left
-    new_dims = [dim_ + 1 if dim_ >= 0 else dim_ for dim_ in dims]
-
-    for dim in new_dims[:-1]:
-        transformed = op.DFT(transformed, axis=dim, inverse=inverse, onesided=False)
-
-    # Torch computers one-sided FFT on the last dimension only.
-    if onesided:
-        transformed = op.DFT(transformed, axis=new_dims[-1], inverse=inverse, onesided=True)
+    """Normalize in forward or backward direction."""
+    # Norm values defined in https://github.com/pytorch/pytorch/blob/758d78790164bfb041555daed380de96e06f78a3/aten/src/ATen/native/SpectralOps.cpp#L117-L131
+    # Norm modes: https://github.com/pytorch/pytorch/blob/758d78790164bfb041555daed380de96e06f78a3/aten/src/ATen/native/SpectralOpsUtils.h#L15-L19
+    # Modes:
+    # 0: no normalization (backward)
+    # 1: "ortho" - divide by 1/sqrt(signal_size) (ortho)
+    # 2: divide by signal_size (forward)
+    signal_size = op.CastLike(signal_size, self)
+    if not inverse:
+        # Forward normalization
+        if normalization == 1:
+            self = op.Div(self, op.Sqrt(signal_size))
+        elif normalization == 2:
+            self = op.Div(self, signal_size)
     else:
-        transformed = op.DFT(transformed, axis=new_dims[-1], inverse=inverse, onesided=False)
-
-    # Remove the batch dimension
-    transformed = op.Squeeze(transformed, axes=[0])
-
-    return _fftn_onnx_normalization(self, transformed, normalization, not inverse, dims)
+        # Backward normalization, accounting for op.DFT already dividing by signal_size
+        if normalization == 0:
+            self = op.Mul(self, signal_size)
+        elif normalization == 1:
+            self = op.Mul(self, op.Sqrt(signal_size))
+    return self
 
 
 @torch_op("aten::_fft_c2c", trace_only=True, complex=True)
@@ -124,14 +59,34 @@ def aten__fft_c2c(
     Standard complex to complex FFT (forward or backward).
     """
 
-    # NOTE: trace_only because we need to negate forward
-    # NOTE: SymInt dim is not support because DFT-17 needs a static axis
-    # TODO(justinchuby): Make dim dynamic and remove trace_only when ONNX provides support
+    # NOTE: SymInt dim is not supported because DFT-17 needs a static axis
 
     # ONNX DFT input assumes the last dimension is the complex dimension.
-    # Thus dim=-1 in PyTorch is dim=-2 in ONNX.
-    dim = [d - 1 if d < 0 else d for d in dim]
-    return _fftn_onnx(self, dim, normalization, inverse=not forward, onesided=False)
+
+    unsqueeze_first_dim = 0 in dim
+    # 1. Add a new dimension for the end and batch dimension, if needed
+    # 2. ONNX DFT input assumes the last dimension is the complex dimension.
+    #       If needed, add 1 to account for the batch dimension.
+
+    if unsqueeze_first_dim:
+        transformed = op.Unsqueeze(self, axes=[0])
+        dim = [d + 1 for d in dim]
+    else:
+        transformed = self
+
+    for dimension in reversed(dim):
+        transformed = op.DFT(transformed, axis=dimension, inverse=not forward, onesided=False)
+        transformed = _fftn_onnx_normalization(
+            transformed,
+            normalization,
+            op.Shape(transformed, start=dimension, end=dimension + 1),
+            not forward,
+        )
+
+    if unsqueeze_first_dim:
+        transformed = op.Squeeze(transformed, axes=[0])
+
+    return transformed
 
 
 @torch_op("aten::_fft_c2r", trace_only=True, complex=True)
@@ -139,24 +94,52 @@ def aten__fft_c2r(
     self: TFloat,
     dim: Sequence[int],
     normalization: int,
-    last_dim_size: INT64,  # pylint: disable=unused-argument
+    last_dim_size: INT64,
 ) -> TFloat:
     """_fft_c2r(Tensor self, int[] dim, int normalization, SymInt last_dim_size) -> Tensor
 
-    Complex to real inverse FFT.
+    Complex to real inverse FFT. Assumes that input tensor is output of previous FFT operation.
     """
-
-    # TODO(justinchuby): Figure out what last_dim_size does
-
-    self_rank = len(self.shape)
-    # ONNX DFT input assumes the last dimension is the complex dimension.
-    # Thus dim=-1 in PyTorch is dim=-2 in ONNX.
-    dim = [(d - 1) + self_rank if d < 0 else d for d in dim]
-    transformed = _fftn_onnx(self, dim, normalization, inverse=True, onesided=False)
-    # Take only the real part
-    real_part = op.Slice(transformed, axes=[-1], starts=[0], ends=[1])
-
-    return op.Squeeze(real_part, axes=[-1])
+    if len(dim) != 1:
+        raise NotImplementedError("Only one dimension is supported for inverse FFT")
+
+    dimension = dim[0]
+    unsqueeze_first_dim = dimension == 0
+    # 1. Add a new dimension for batch dimension, if needed
+    # 2. ONNX DFT input assumes the last dimension is the complex dimension.
+    #       If needed, add 1 to account for the batch dimension.
+
+    if unsqueeze_first_dim:
+        transformed = op.Unsqueeze(self, axes=[0])
+        dimension = 1
+    else:
+        transformed = self
+
+    # Torch truncates/pads on the last dimension only. Typically, the only valid values that can be passed
+    # into PyTorch are n or n//2+1, where n is self.shape[dim[-1]], but this is not always the case, so we
+    # place no such restriction on the ONNX side.
+    transformed = op.DFT(
+        transformed,
+        dft_length=last_dim_size,
+        axis=dimension,
+        inverse=True,
+        onesided=False,
+    )
+    transformed = _fftn_onnx_normalization(
+        transformed,
+        normalization,
+        op.Shape(transformed, start=dimension, end=dimension + 1),
+        inverse=True,
+    )
+
+    if unsqueeze_first_dim:
+        transformed = op.Squeeze(transformed, axes=[0])
+
+    # Remove the imaginary part
+    transformed = op.Slice(transformed, [0], [1], [-1])
+    transformed = op.Squeeze(transformed, axes=[-1])
+
+    return transformed
 
 
 @torch_op("aten::_fft_r2c", trace_only=True)
@@ -168,17 +151,37 @@ def aten__fft_r2c(
     Real to complex forward FFT.
     """
 
-    # Add a new dimension at the end
-    signal = op.Unsqueeze(self, axes=[-1])
     # No need to fill the imaginary part because ONNX DFT accepts real inputs
     # https://onnx.ai/onnx/operators/onnx__DFT.html#inputs
 
-    self_rank = len(self.shape)
-    # ONNX DFT input assumes the last dimension is the complex dimension.
-    # Thus dim=-1 in PyTorch is dim=-2 in ONNX.
-    dim = [(d - 1) + self_rank if d < 0 else d for d in dim]
+    unsqueeze_first_dim = 0 in dim
+    # 1. Add a new dimension for the end and batch dimension, if needed
+    # 2. ONNX DFT input assumes the last dimension is the complex dimension.
+    #       If needed, add 1 to account for the batch dimension.
+
+    if unsqueeze_first_dim:
+        transformed = op.Unsqueeze(self, axes=[0, -1])
+        dim = [d + 1 for d in dim]
+    else:
+        transformed = op.Unsqueeze(self, axes=[-1])
+
+    for idx, dimension in enumerate(reversed(dim)):
+        transformed = _fftn_onnx_normalization(
+            transformed,
+            normalization,
+            op.Shape(transformed, start=dimension, end=dimension + 1),
+            inverse=False,
+        )
+        if idx > 0:
+            transformed = op.DFT(transformed, axis=dimension, inverse=False, onesided=False)
+        else:
+            # Torch computes one-sided FFT on the last dimension only.
+            transformed = op.DFT(transformed, axis=dimension, inverse=False, onesided=onesided)
+
+    if unsqueeze_first_dim:
+        transformed = op.Squeeze(transformed, axes=[0])
 
-    return _fftn_onnx(signal, dim, normalization, inverse=False, onesided=onesided)
+    return transformed
 
 
 def aten_fft_fft(
diff --git a/onnxscript/ir/tensor_adapters_test.py b/onnxscript/ir/tensor_adapters_test.py
index 34034ac51f..b9ba55be57 100644
--- a/onnxscript/ir/tensor_adapters_test.py
+++ b/onnxscript/ir/tensor_adapters_test.py
@@ -54,25 +54,25 @@ def test_numpy_returns_correct_dtype(self, dtype: torch.dtype, np_dtype):
 
     @parameterized.parameterized.expand(
         [
-            (torch.bfloat16),
-            (torch.bool),
-            (torch.complex128),
-            (torch.complex64),
-            (torch.float16),
-            (torch.float32),
-            (torch.float64),
-            (torch.float8_e4m3fn),
-            (torch.float8_e4m3fnuz),
-            (torch.float8_e5m2),
-            (torch.float8_e5m2fnuz),
-            (torch.int16),
-            (torch.int32),
-            (torch.int64),
-            (torch.int8),
-            (torch.uint16),
-            (torch.uint32),
-            (torch.uint64),
-            (torch.uint8),
+            (torch.bfloat16,),
+            (torch.bool,),
+            (torch.complex128,),
+            (torch.complex64,),
+            (torch.float16,),
+            (torch.float32,),
+            (torch.float64,),
+            (torch.float8_e4m3fn,),
+            (torch.float8_e4m3fnuz,),
+            (torch.float8_e5m2,),
+            (torch.float8_e5m2fnuz,),
+            (torch.int16,),
+            (torch.int32,),
+            (torch.int64,),
+            (torch.int8,),
+            (torch.uint16,),
+            (torch.uint32,),
+            (torch.uint64,),
+            (torch.uint8,),
         ],
     )
     def test_tobytes(self, dtype: torch.dtype):
diff --git a/tests/function_libs/torch_lib/extra_opinfo.py b/tests/function_libs/torch_lib/extra_opinfo.py
index 70a1e0547f..26b75bf93b 100644
--- a/tests/function_libs/torch_lib/extra_opinfo.py
+++ b/tests/function_libs/torch_lib/extra_opinfo.py
@@ -684,24 +684,38 @@ def sample_inputs__fft_r2c(self, device, dtype, requires_grad=False, **_):
 
 def sample_inputs__fft_c2r(self, device, dtype, requires_grad=False, **_):
     del self  # Unused
-    oned_tensor, nd_tensor = _prepare_data_for_fft_ops(device, dtype, requires_grad)
-
+    real_dtype = torch.float
+    if dtype == torch.complex128:
+        real_dtype = torch.double
+    oned_tensor, nd_tensor = _prepare_data_for_fft_ops(device, real_dtype, requires_grad)
+    oned_tensor_result = oned_tensor()
+    nd_tensor_result = nd_tensor()
+    complex_oned_tensor = torch.ops.aten._fft_r2c.default(  # pylint: disable=protected-access
+        oned_tensor_result, [0], normalization=0, onesided=False
+    )
+    # for normalization in (0, 1, 2):
     for normalization in (0, 1, 2):
         # 1-D
         yield opinfo_core.SampleInput(
-            oned_tensor(), dim=(0,), normalization=normalization, last_dim_size=12
+            complex_oned_tensor,
+            dim=(0,),
+            normalization=normalization,
+            last_dim_size=31,
         )
         # N-D
         for dim in [
             (0,),
             (1,),
             (2,),
-            (1, 2),
-            (0, 1),
-            (0, 1, 2),
         ]:
+            complex_nd_tensor = torch.ops.aten._fft_r2c.default(  # pylint: disable=protected-access
+                nd_tensor_result, dim, normalization=0, onesided=False
+            )
             yield opinfo_core.SampleInput(
-                nd_tensor(), dim=dim, normalization=normalization, last_dim_size=6
+                complex_nd_tensor,
+                dim=dim,
+                normalization=normalization,
+                last_dim_size=complex_nd_tensor.shape[dim[-1]],
             )
 
 
diff --git a/tests/function_libs/torch_lib/ops_test_data.py b/tests/function_libs/torch_lib/ops_test_data.py
index 54e1e8cceb..3628ed8c45 100644
--- a/tests/function_libs/torch_lib/ops_test_data.py
+++ b/tests/function_libs/torch_lib/ops_test_data.py
@@ -452,9 +452,6 @@ def _where_input_wrangler(
         fft_ops.aten__fft_c2r,
         tolerance={torch.complex64: (3e-3, 1.8e-4)},
         complex=True,
-    ).xfail(
-        dtypes=(torch.complex64,),
-        reason="fixme: the result is wrong: https://github.com/microsoft/onnxscript/pull/926",
     ),
     TorchLibOpInfo(
         "ops.aten._fft_r2c",  # Custom from extra_opinfo


Reviews:
user: login: github-advanced-security[bot]
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-03-27T17:00:29Z
user: login: github-advanced-security[bot]
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-03-27T17:02:41Z
user: login: github-advanced-security[bot]
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-03-27T17:18:23Z
user: login: github-advanced-security[bot]
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-03-27T17:29:21Z
user: login: gramalingam
state: COMMENTED
author_association: COLLABORATOR
submitted_at: 2025-03-27T19:28:32Z
user: login: bmehta001
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-03-27T20:29:34Z
user: login: titaiwangms
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-03-27T21:54:40Z
user: login: titaiwangms
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-03-27T21:56:25Z
user: login: justinchuby
state: COMMENTED
author_association: COLLABORATOR
submitted_at: 2025-03-28T02:06:26Z
user: login: justinchuby
state: COMMENTED
author_association: COLLABORATOR
submitted_at: 2025-03-31T16:24:43Z
user: login: titaiwangms
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-03-31T16:28:40Z
user: login: github-advanced-security[bot]
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-04-10T17:28:33Z
user: login: github-advanced-security[bot]
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-04-11T01:11:01Z
user: login: justinchuby
state: COMMENTED
author_association: COLLABORATOR
submitted_at: 2025-04-11T01:29:04Z
user: login: justinchuby
state: COMMENTED
author_association: COLLABORATOR
submitted_at: 2025-04-11T01:31:44Z
user: login: justinchuby
state: COMMENTED
author_association: COLLABORATOR
submitted_at: 2025-04-11T15:12:02Z
user: login: justinchuby
state: COMMENTED
author_association: COLLABORATOR
submitted_at: 2025-04-11T15:13:19Z
user: login: justinchuby
state: COMMENTED
author_association: COLLABORATOR
submitted_at: 2025-04-13T14:43:37Z
user: login: bmehta001
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-04-14T08:27:37Z
user: login: bmehta001
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-04-14T08:28:23Z
user: login: bmehta001
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-04-14T08:28:32Z
user: login: bmehta001
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-04-14T08:31:08Z
user: login: justinchuby
state: COMMENTED
author_association: COLLABORATOR
submitted_at: 2025-04-14T14:56:52Z
user: login: justinchuby
state: APPROVED
author_association: COLLABORATOR
submitted_at: 2025-04-16T01:28:35Z
user: login: justinchuby
state: COMMENTED
author_association: COLLABORATOR
submitted_at: 2025-04-16T01:34:26Z
user: login: justinchuby
state: APPROVED
author_association: COLLABORATOR
submitted_at: 2025-04-16T01:34:29Z
user: login: justinchuby
state: APPROVED
author_association: COLLABORATOR
submitted_at: 2025-04-16T01:34:34Z

Review comments:
diff_hunk: @@ -12,108 +10,53 @@
 
 from __future__ import annotations
 
-from typing import Optional, Sequence
+from typing import Literal, Optional, Sequence
path: onnxscript/function_libs/torch_lib/ops/fft.py
user: login: github-advanced-security[bot]
body: ## Unused import

Import of 'Literal' is not used.

[Show more details](https://github.com/microsoft/onnxscript/security/code-scanning/16123)
created_at: 2025-03-27T17:00:29Z
author_association: CONTRIBUTOR
diff_hunk: @@ -124,39 +67,74 @@
     Standard complex to complex FFT (forward or backward).
     """
 
-    # NOTE: trace_only because we need to negate forward
-    # NOTE: SymInt dim is not support because DFT-17 needs a static axis
-    # TODO(justinchuby): Make dim dynamic and remove trace_only when ONNX provides support
+    # NOTE: SymInt dim is not supported because DFT-17 needs a static axis
 
     # ONNX DFT input assumes the last dimension is the complex dimension.
     # Thus dim=-1 in PyTorch is dim=-2 in ONNX.
-    dim = [d - 1 if d < 0 else d for d in dim]
-    return _fftn_onnx(self, dim, normalization, inverse=not forward, onesided=False)
+    assert(dim[2] in dim == 2, "Unexpected input size")
path: onnxscript/function_libs/torch_lib/ops/fft.py
user: login: github-advanced-security[bot]
body: ## Asserting a tuple

Assertion of non-empty tuple is always True.

[Show more details](https://github.com/microsoft/onnxscript/security/code-scanning/16506)
created_at: 2025-03-27T17:00:29Z
author_association: CONTRIBUTOR
diff_hunk: @@ -124,39 +67,74 @@
     Standard complex to complex FFT (forward or backward).
     """
 
-    # NOTE: trace_only because we need to negate forward
-    # NOTE: SymInt dim is not support because DFT-17 needs a static axis
-    # TODO(justinchuby): Make dim dynamic and remove trace_only when ONNX provides support
+    # NOTE: SymInt dim is not supported because DFT-17 needs a static axis
 
     # ONNX DFT input assumes the last dimension is the complex dimension.
     # Thus dim=-1 in PyTorch is dim=-2 in ONNX.
-    dim = [d - 1 if d < 0 else d for d in dim]
-    return _fftn_onnx(self, dim, normalization, inverse=not forward, onesided=False)
+    assert(dim[2] in dim == 2, "Unexpected input size")
+
+    signal = self
+    self_rank = len(self.shape)
+    signal_size = op.Size(signal)
+
+    # ONNX DFT input assumes the last dimension is the complex dimension.
+    # Thus dim=-1 in PyTorch is dim=-2 in ONNX.
+    dim = [(d - 1) + self_rank if d < 0 else d for d in dim]
+
+    transformed = signal
+
+    for dimension in reversed(dim):
+        transformed = op.DFT(transformed, axis=dimension, inverse=not forward, onesided=False)
+        if forward:
+            transformed = _fftn_onnx_normalization(transformed, normalization, signal_size)
+        else:
+            transformed = _fftn_onnx_inverse_normalization(transformed, normalization, signal_size)
+
+    # Unsure if output format is correct
+    return transformed
 
 
 @torch_op("aten::_fft_c2r", trace_only=True, complex=True)
 def aten__fft_c2r(
     self: TFloat,
     dim: Sequence[int],
     normalization: int,
-    last_dim_size: INT64,  # pylint: disable=unused-argument
+    last_dim_size: INT64,
 ) -> TFloat:
     """_fft_c2r(Tensor self, int[] dim, int normalization, SymInt last_dim_size) -> Tensor
 
     Complex to real inverse FFT.
     """
+    assert(dim[2] in dim == 2, "Unexpected input size")
path: onnxscript/function_libs/torch_lib/ops/fft.py
user: login: github-advanced-security[bot]
body: ## Asserting a tuple

Assertion of non-empty tuple is always True.

[Show more details](https://github.com/microsoft/onnxscript/security/code-scanning/16507)
created_at: 2025-03-27T17:00:29Z
author_association: CONTRIBUTOR
diff_hunk: @@ -1,7 +1,5 @@
-# --------------------------------------------------------------------------
-# Copyright (c) Microsoft Corporation. All rights reserved.
+# Copyright (c) Microsoft Corporation.
path: onnxscript/function_libs/torch_lib/ops/fft.py
user: login: github-advanced-security[bot]
body: ## RUFF-FORMAT/format

Run `lintrunner -a` to apply this patch.

[Show more details](https://github.com/microsoft/onnxscript/security/code-scanning/12483)
created_at: 2025-03-27T17:02:40Z
author_association: CONTRIBUTOR
diff_hunk: @@ -1,7 +1,5 @@
-# --------------------------------------------------------------------------
-# Copyright (c) Microsoft Corporation. All rights reserved.
+# Copyright (c) Microsoft Corporation.
path: onnxscript/function_libs/torch_lib/ops/fft.py
user: login: github-advanced-security[bot]
body: ## RUFF/format

Run `lintrunner -a` to apply this patch.

[Show more details](https://github.com/microsoft/onnxscript/security/code-scanning/12489)
created_at: 2025-03-27T17:02:40Z
author_association: CONTRIBUTOR
diff_hunk: @@ -12,108 +10,53 @@
 
 from __future__ import annotations
 
-from typing import Optional, Sequence
+from typing import Literal, Optional, Sequence
path: onnxscript/function_libs/torch_lib/ops/fft.py
user: login: github-advanced-security[bot]
body: ## PYLINT/W0611

Unused Literal imported from typing (unused-import)
See [unused-import](https://pylint.pycqa.org/en/latest/user_guide/messages/warning/unused-import.html).

To disable, use `  # pylint: disable=unused-import`

[Show more details](https://github.com/microsoft/onnxscript/security/code-scanning/12500)
created_at: 2025-03-27T17:02:40Z
author_association: CONTRIBUTOR
diff_hunk: @@ -12,108 +10,53 @@
 
 from __future__ import annotations
 
-from typing import Optional, Sequence
+from typing import Literal, Optional, Sequence
path: onnxscript/function_libs/torch_lib/ops/fft.py
user: login: github-advanced-security[bot]
body: ## RUFF/F401

`typing.Literal` imported but unused.
See https://docs.astral.sh/ruff/rules/unused-import

[Show more details](https://github.com/microsoft/onnxscript/security/code-scanning/12485)
created_at: 2025-03-27T17:02:40Z
author_association: CONTRIBUTOR
diff_hunk: @@ -12,108 +10,53 @@
 
 from __future__ import annotations
 
-from typing import Optional, Sequence
+from typing import Literal, Optional, Sequence
 
 from onnxscript import INT64
 from onnxscript.function_libs.torch_lib.registration import torch_op
 from onnxscript.function_libs.torch_lib.tensor_typing import TFloat
 from onnxscript.onnx_opset import opset18 as op
 from onnxscript.onnx_types import TensorType
 
-
-@torch_op(
-    ("aten::_fft_c2c", "aten::_fft_c2r", "aten::_fft_r2c"),
-    private=True,
-    complex=True,
-    trace_only=True,
-)
 def _fftn_onnx_normalization(
-    self,
-    transformed: TFloat,
+    self: TFloat,
     normalization: int,
-    forward: bool,
-    dims: Sequence[int],
+    signal_size: INT64,
 ) -> TFloat:
-    # Obtain the total_sample_count (n) for normalization
-    self_shape = op.Shape(self)
-    total_sample_count = op.ReduceProd(op.Gather(self_shape, dims), keepdims=0)
-    total_sample_count = op.CastLike(total_sample_count, transformed)
-
-    # Normalize the result
-    # Reference https://pytorch.org/docs/stable/generated/torch.fft.fftn.html#torch.fft.fftn
-    # Reference https://github.com/pytorch/pytorch/blob/d090c18fcaaba6e1b5cb474a89058cf6081c8275/torch/_refs/fft.py#L42
+    """
path: onnxscript/function_libs/torch_lib/ops/fft.py
user: login: github-advanced-security[bot]
body: ## RUFF/D419

Docstring is empty.
See https://docs.astral.sh/ruff/rules/empty-docstring

[Show more details](https://github.com/microsoft/onnxscript/security/code-scanning/16509)
created_at: 2025-03-27T17:02:40Z
author_association: CONTRIBUTOR
diff_hunk: @@ -124,39 +67,74 @@
     Standard complex to complex FFT (forward or backward).
     """
 
-    # NOTE: trace_only because we need to negate forward
-    # NOTE: SymInt dim is not support because DFT-17 needs a static axis
-    # TODO(justinchuby): Make dim dynamic and remove trace_only when ONNX provides support
+    # NOTE: SymInt dim is not supported because DFT-17 needs a static axis
 
     # ONNX DFT input assumes the last dimension is the complex dimension.
     # Thus dim=-1 in PyTorch is dim=-2 in ONNX.
-    dim = [d - 1 if d < 0 else d for d in dim]
-    return _fftn_onnx(self, dim, normalization, inverse=not forward, onesided=False)
+    assert(dim[2] in dim == 2, "Unexpected input size")
path: onnxscript/function_libs/torch_lib/ops/fft.py
user: login: github-advanced-security[bot]
body: ## PYLINT/W0199

Assert called on a populated tuple. Did you mean 'assert x,y'? (assert-on-tuple)
See [assert-on-tuple](https://pylint.pycqa.org/en/latest/user_guide/messages/warning/assert-on-tuple.html).

To disable, use `  # pylint: disable=assert-on-tuple`

[Show more details](https://github.com/microsoft/onnxscript/security/code-scanning/16513)
created_at: 2025-03-27T17:02:40Z
author_association: CONTRIBUTOR
diff_hunk: @@ -124,39 +67,74 @@
     Standard complex to complex FFT (forward or backward).
     """
 
-    # NOTE: trace_only because we need to negate forward
-    # NOTE: SymInt dim is not support because DFT-17 needs a static axis
-    # TODO(justinchuby): Make dim dynamic and remove trace_only when ONNX provides support
+    # NOTE: SymInt dim is not supported because DFT-17 needs a static axis
 
     # ONNX DFT input assumes the last dimension is the complex dimension.
     # Thus dim=-1 in PyTorch is dim=-2 in ONNX.
-    dim = [d - 1 if d < 0 else d for d in dim]
-    return _fftn_onnx(self, dim, normalization, inverse=not forward, onesided=False)
+    assert(dim[2] in dim == 2, "Unexpected input size")
path: onnxscript/function_libs/torch_lib/ops/fft.py
user: login: github-advanced-security[bot]
body: ## RUFF/F631

Assert test is a non-empty tuple, which is always `True`.
See https://docs.astral.sh/ruff/rules/assert-tuple

[Show more details](https://github.com/microsoft/onnxscript/security/code-scanning/16511)
created_at: 2025-03-27T17:02:40Z
author_association: CONTRIBUTOR
diff_hunk: @@ -124,39 +67,74 @@
     Standard complex to complex FFT (forward or backward).
     """
 
-    # NOTE: trace_only because we need to negate forward
-    # NOTE: SymInt dim is not support because DFT-17 needs a static axis
-    # TODO(justinchuby): Make dim dynamic and remove trace_only when ONNX provides support
+    # NOTE: SymInt dim is not supported because DFT-17 needs a static axis
 
     # ONNX DFT input assumes the last dimension is the complex dimension.
     # Thus dim=-1 in PyTorch is dim=-2 in ONNX.
-    dim = [d - 1 if d < 0 else d for d in dim]
-    return _fftn_onnx(self, dim, normalization, inverse=not forward, onesided=False)
+    assert(dim[2] in dim == 2, "Unexpected input size")
path: onnxscript/function_libs/torch_lib/ops/fft.py
user: login: github-advanced-security[bot]
body: ## PYLINT/W3601

Suspicious 2-part chained comparison using semantically incompatible operators ('==' and 'in') (bad-chained-comparison)
See [bad-chained-comparison](https://pylint.pycqa.org/en/latest/user_guide/messages/warning/bad-chained-comparison.html).

To disable, use `  # pylint: disable=bad-chained-comparison`

[Show more details](https://github.com/microsoft/onnxscript/security/code-scanning/16514)
created_at: 2025-03-27T17:02:41Z
author_association: CONTRIBUTOR
diff_hunk: @@ -124,39 +67,74 @@
     Standard complex to complex FFT (forward or backward).
     """
 
-    # NOTE: trace_only because we need to negate forward
-    # NOTE: SymInt dim is not support because DFT-17 needs a static axis
-    # TODO(justinchuby): Make dim dynamic and remove trace_only when ONNX provides support
+    # NOTE: SymInt dim is not supported because DFT-17 needs a static axis
 
     # ONNX DFT input assumes the last dimension is the complex dimension.
     # Thus dim=-1 in PyTorch is dim=-2 in ONNX.
-    dim = [d - 1 if d < 0 else d for d in dim]
-    return _fftn_onnx(self, dim, normalization, inverse=not forward, onesided=False)
+    assert(dim[2] in dim == 2, "Unexpected input size")
+
+    signal = self
+    self_rank = len(self.shape)
+    signal_size = op.Size(signal)
+
+    # ONNX DFT input assumes the last dimension is the complex dimension.
+    # Thus dim=-1 in PyTorch is dim=-2 in ONNX.
+    dim = [(d - 1) + self_rank if d < 0 else d for d in dim]
+
+    transformed = signal
+
+    for dimension in reversed(dim):
+        transformed = op.DFT(transformed, axis=dimension, inverse=not forward, onesided=False)
+        if forward:
+            transformed = _fftn_onnx_normalization(transformed, normalization, signal_size)
+        else:
+            transformed = _fftn_onnx_inverse_normalization(transformed, normalization, signal_size)
+
+    # Unsure if output format is correct
+    return transformed
 
 
 @torch_op("aten::_fft_c2r", trace_only=True, complex=True)
 def aten__fft_c2r(
     self: TFloat,
     dim: Sequence[int],
     normalization: int,
-    last_dim_size: INT64,  # pylint: disable=unused-argument
+    last_dim_size: INT64,
 ) -> TFloat:
     """_fft_c2r(Tensor self, int[] dim, int normalization, SymInt last_dim_size) -> Tensor
 
     Complex to real inverse FFT.
     """
+    assert(dim[2] in dim == 2, "Unexpected input size")
path: onnxscript/function_libs/torch_lib/ops/fft.py
user: login: github-advanced-security[bot]
body: ## PYLINT/W0199

Assert called on a populated tuple. Did you mean 'assert x,y'? (assert-on-tuple)
See [assert-on-tuple](https://pylint.pycqa.org/en/latest/user_guide/messages/warning/assert-on-tuple.html).

To disable, use `  # pylint: disable=assert-on-tuple`

[Show more details](https://github.com/microsoft/onnxscript/security/code-scanning/16515)
created_at: 2025-03-27T17:02:41Z
author_association: CONTRIBUTOR
diff_hunk: @@ -124,39 +67,74 @@
     Standard complex to complex FFT (forward or backward).
     """
 
-    # NOTE: trace_only because we need to negate forward
-    # NOTE: SymInt dim is not support because DFT-17 needs a static axis
-    # TODO(justinchuby): Make dim dynamic and remove trace_only when ONNX provides support
+    # NOTE: SymInt dim is not supported because DFT-17 needs a static axis
 
     # ONNX DFT input assumes the last dimension is the complex dimension.
     # Thus dim=-1 in PyTorch is dim=-2 in ONNX.
-    dim = [d - 1 if d < 0 else d for d in dim]
-    return _fftn_onnx(self, dim, normalization, inverse=not forward, onesided=False)
+    assert(dim[2] in dim == 2, "Unexpected input size")
+
+    signal = self
+    self_rank = len(self.shape)
+    signal_size = op.Size(signal)
+
+    # ONNX DFT input assumes the last dimension is the complex dimension.
+    # Thus dim=-1 in PyTorch is dim=-2 in ONNX.
+    dim = [(d - 1) + self_rank if d < 0 else d for d in dim]
+
+    transformed = signal
+
+    for dimension in reversed(dim):
+        transformed = op.DFT(transformed, axis=dimension, inverse=not forward, onesided=False)
+        if forward:
+            transformed = _fftn_onnx_normalization(transformed, normalization, signal_size)
+        else:
+            transformed = _fftn_onnx_inverse_normalization(transformed, normalization, signal_size)
+
+    # Unsure if output format is correct
+    return transformed
 
 
 @torch_op("aten::_fft_c2r", trace_only=True, complex=True)
 def aten__fft_c2r(
     self: TFloat,
     dim: Sequence[int],
     normalization: int,
-    last_dim_size: INT64,  # pylint: disable=unused-argument
+    last_dim_size: INT64,
 ) -> TFloat:
     """_fft_c2r(Tensor self, int[] dim, int normalization, SymInt last_dim_size) -> Tensor
 
     Complex to real inverse FFT.
     """
+    assert(dim[2] in dim == 2, "Unexpected input size")
path: onnxscript/function_libs/torch_lib/ops/fft.py
user: login: github-advanced-security[bot]
body: ## RUFF/F631

Assert test is a non-empty tuple, which is always `True`.
See https://docs.astral.sh/ruff/rules/assert-tuple

[Show more details](https://github.com/microsoft/onnxscript/security/code-scanning/16512)
created_at: 2025-03-27T17:02:41Z
author_association: CONTRIBUTOR
diff_hunk: @@ -124,39 +67,74 @@
     Standard complex to complex FFT (forward or backward).
     """
 
-    # NOTE: trace_only because we need to negate forward
-    # NOTE: SymInt dim is not support because DFT-17 needs a static axis
-    # TODO(justinchuby): Make dim dynamic and remove trace_only when ONNX provides support
+    # NOTE: SymInt dim is not supported because DFT-17 needs a static axis
 
     # ONNX DFT input assumes the last dimension is the complex dimension.
     # Thus dim=-1 in PyTorch is dim=-2 in ONNX.
-    dim = [d - 1 if d < 0 else d for d in dim]
-    return _fftn_onnx(self, dim, normalization, inverse=not forward, onesided=False)
+    assert(dim[2] in dim == 2, "Unexpected input size")
+
+    signal = self
+    self_rank = len(self.shape)
+    signal_size = op.Size(signal)
+
+    # ONNX DFT input assumes the last dimension is the complex dimension.
+    # Thus dim=-1 in PyTorch is dim=-2 in ONNX.
+    dim = [(d - 1) + self_rank if d < 0 else d for d in dim]
+
+    transformed = signal
+
+    for dimension in reversed(dim):
+        transformed = op.DFT(transformed, axis=dimension, inverse=not forward, onesided=False)
+        if forward:
+            transformed = _fftn_onnx_normalization(transformed, normalization, signal_size)
+        else:
+            transformed = _fftn_onnx_inverse_normalization(transformed, normalization, signal_size)
+
+    # Unsure if output format is correct
+    return transformed
 
 
 @torch_op("aten::_fft_c2r", trace_only=True, complex=True)
 def aten__fft_c2r(
     self: TFloat,
     dim: Sequence[int],
     normalization: int,
-    last_dim_size: INT64,  # pylint: disable=unused-argument
+    last_dim_size: INT64,
 ) -> TFloat:
     """_fft_c2r(Tensor self, int[] dim, int normalization, SymInt last_dim_size) -> Tensor
 
     Complex to real inverse FFT.
     """
+    assert(dim[2] in dim == 2, "Unexpected input size")
path: onnxscript/function_libs/torch_lib/ops/fft.py
user: login: github-advanced-security[bot]
body: ## PYLINT/W3601

Suspicious 2-part chained comparison using semantically incompatible operators ('==' and 'in') (bad-chained-comparison)
See [bad-chained-comparison](https://pylint.pycqa.org/en/latest/user_guide/messages/warning/bad-chained-comparison.html).

To disable, use `  # pylint: disable=bad-chained-comparison`

[Show more details](https://github.com/microsoft/onnxscript/security/code-scanning/16516)
created_at: 2025-03-27T17:02:41Z
author_association: CONTRIBUTOR
diff_hunk: @@ -20,100 +18,45 @@
 from onnxscript.onnx_opset import opset18 as op
 from onnxscript.onnx_types import TensorType
 
-
-@torch_op(
-    ("aten::_fft_c2c", "aten::_fft_c2r", "aten::_fft_r2c"),
-    private=True,
-    complex=True,
-    trace_only=True,
-)
 def _fftn_onnx_normalization(
-    self,
-    transformed: TFloat,
+    self: TFloat,
     normalization: int,
-    forward: bool,
-    dims: Sequence[int],
+    signal_size: INT64,
 ) -> TFloat:
-    # Obtain the total_sample_count (n) for normalization
-    self_shape = op.Shape(self)
-    total_sample_count = op.ReduceProd(op.Gather(self_shape, dims), keepdims=0)
-    total_sample_count = op.CastLike(total_sample_count, transformed)
-
-    # Normalize the result
-    # Reference https://pytorch.org/docs/stable/generated/torch.fft.fftn.html#torch.fft.fftn
-    # Reference https://github.com/pytorch/pytorch/blob/d090c18fcaaba6e1b5cb474a89058cf6081c8275/torch/_refs/fft.py#L42
+    """
path: onnxscript/function_libs/torch_lib/ops/fft.py
user: login: github-advanced-security[bot]
body: ## RUFF/D419

Docstring is empty.
See https://docs.astral.sh/ruff/rules/empty-docstring

[Show more details](https://github.com/microsoft/onnxscript/security/code-scanning/16518)
created_at: 2025-03-27T17:18:22Z
author_association: CONTRIBUTOR
diff_hunk: @@ -124,39 +67,74 @@
     Standard complex to complex FFT (forward or backward).
     """
 
-    # NOTE: trace_only because we need to negate forward
-    # NOTE: SymInt dim is not support because DFT-17 needs a static axis
-    # TODO(justinchuby): Make dim dynamic and remove trace_only when ONNX provides support
+    # NOTE: SymInt dim is not supported because DFT-17 needs a static axis
 
     # ONNX DFT input assumes the last dimension is the complex dimension.
     # Thus dim=-1 in PyTorch is dim=-2 in ONNX.
-    dim = [d - 1 if d < 0 else d for d in dim]
-    return _fftn_onnx(self, dim, normalization, inverse=not forward, onesided=False)
+    assert(dim[-1] in dim == 2, "Unexpected input size")
path: onnxscript/function_libs/torch_lib/ops/fft.py
user: login: github-advanced-security[bot]
body: ## PYLINT/W0199

Assert called on a populated tuple. Did you mean 'assert x,y'? (assert-on-tuple)
See [assert-on-tuple](https://pylint.pycqa.org/en/latest/user_guide/messages/warning/assert-on-tuple.html).

To disable, use `  # pylint: disable=assert-on-tuple`

[Show more details](https://github.com/microsoft/onnxscript/security/code-scanning/16522)
created_at: 2025-03-27T17:18:23Z
author_association: CONTRIBUTOR
diff_hunk: @@ -124,39 +67,74 @@
     Standard complex to complex FFT (forward or backward).
     """
 
-    # NOTE: trace_only because we need to negate forward
-    # NOTE: SymInt dim is not support because DFT-17 needs a static axis
-    # TODO(justinchuby): Make dim dynamic and remove trace_only when ONNX provides support
+    # NOTE: SymInt dim is not supported because DFT-17 needs a static axis
 
     # ONNX DFT input assumes the last dimension is the complex dimension.
     # Thus dim=-1 in PyTorch is dim=-2 in ONNX.
-    dim = [d - 1 if d < 0 else d for d in dim]
-    return _fftn_onnx(self, dim, normalization, inverse=not forward, onesided=False)
+    assert(dim[-1] in dim == 2, "Unexpected input size")
path: onnxscript/function_libs/torch_lib/ops/fft.py
user: login: github-advanced-security[bot]
body: ## RUFF/F631

Assert test is a non-empty tuple, which is always `True`.
See https://docs.astral.sh/ruff/rules/assert-tuple

[Show more details](https://github.com/microsoft/onnxscript/security/code-scanning/16520)
created_at: 2025-03-27T17:18:23Z
author_association: CONTRIBUTOR
diff_hunk: @@ -124,39 +67,74 @@
     Standard complex to complex FFT (forward or backward).
     """
 
-    # NOTE: trace_only because we need to negate forward
-    # NOTE: SymInt dim is not support because DFT-17 needs a static axis
-    # TODO(justinchuby): Make dim dynamic and remove trace_only when ONNX provides support
+    # NOTE: SymInt dim is not supported because DFT-17 needs a static axis
 
     # ONNX DFT input assumes the last dimension is the complex dimension.
     # Thus dim=-1 in PyTorch is dim=-2 in ONNX.
-    dim = [d - 1 if d < 0 else d for d in dim]
-    return _fftn_onnx(self, dim, normalization, inverse=not forward, onesided=False)
+    assert(dim[-1] in dim == 2, "Unexpected input size")
path: onnxscript/function_libs/torch_lib/ops/fft.py
user: login: github-advanced-security[bot]
body: ## PYLINT/W3601

Suspicious 2-part chained comparison using semantically incompatible operators ('==' and 'in') (bad-chained-comparison)
See [bad-chained-comparison](https://pylint.pycqa.org/en/latest/user_guide/messages/warning/bad-chained-comparison.html).

To disable, use `  # pylint: disable=bad-chained-comparison`

[Show more details](https://github.com/microsoft/onnxscript/security/code-scanning/16523)
created_at: 2025-03-27T17:18:23Z
author_association: CONTRIBUTOR
diff_hunk: @@ -124,39 +67,74 @@
     Standard complex to complex FFT (forward or backward).
     """
 
-    # NOTE: trace_only because we need to negate forward
-    # NOTE: SymInt dim is not support because DFT-17 needs a static axis
-    # TODO(justinchuby): Make dim dynamic and remove trace_only when ONNX provides support
+    # NOTE: SymInt dim is not supported because DFT-17 needs a static axis
 
     # ONNX DFT input assumes the last dimension is the complex dimension.
     # Thus dim=-1 in PyTorch is dim=-2 in ONNX.
-    dim = [d - 1 if d < 0 else d for d in dim]
-    return _fftn_onnx(self, dim, normalization, inverse=not forward, onesided=False)
+    assert(dim[-1] in dim == 2, "Unexpected input size")
+
+    signal = self
+    self_rank = len(self.shape)
+    signal_size = op.Size(signal)
+
+    # ONNX DFT input assumes the last dimension is the complex dimension.
+    # Thus dim=-1 in PyTorch is dim=-2 in ONNX.
+    dim = [(d - 1) + self_rank if d < 0 else d for d in dim]
+
+    transformed = signal
+
+    for dimension in reversed(dim):
+        transformed = op.DFT(transformed, axis=dimension, inverse=not forward, onesided=False)
+        if forward:
+            transformed = _fftn_onnx_normalization(transformed, normalization, signal_size)
+        else:
+            transformed = _fftn_onnx_inverse_normalization(transformed, normalization, signal_size)
+
+    # Unsure if output format is correct
+    return transformed
 
 
 @torch_op("aten::_fft_c2r", trace_only=True, complex=True)
 def aten__fft_c2r(
     self: TFloat,
     dim: Sequence[int],
     normalization: int,
-    last_dim_size: INT64,  # pylint: disable=unused-argument
+    last_dim_size: INT64,
 ) -> TFloat:
     """_fft_c2r(Tensor self, int[] dim, int normalization, SymInt last_dim_size) -> Tensor
 
     Complex to real inverse FFT.
     """
+    assert(dim[-1] in dim == 2, "Unexpected input size")
path: onnxscript/function_libs/torch_lib/ops/fft.py
user: login: github-advanced-security[bot]
body: ## PYLINT/W0199

Assert called on a populated tuple. Did you mean 'assert x,y'? (assert-on-tuple)
See [assert-on-tuple](https://pylint.pycqa.org/en/latest/user_guide/messages/warning/assert-on-tuple.html).

To disable, use `  # pylint: disable=assert-on-tuple`

[Show more details](https://github.com/microsoft/onnxscript/security/code-scanning/16524)
created_at: 2025-03-27T17:18:23Z
author_association: CONTRIBUTOR
diff_hunk: @@ -124,39 +67,74 @@
     Standard complex to complex FFT (forward or backward).
     """
 
-    # NOTE: trace_only because we need to negate forward
-    # NOTE: SymInt dim is not support because DFT-17 needs a static axis
-    # TODO(justinchuby): Make dim dynamic and remove trace_only when ONNX provides support
+    # NOTE: SymInt dim is not supported because DFT-17 needs a static axis
 
     # ONNX DFT input assumes the last dimension is the complex dimension.
     # Thus dim=-1 in PyTorch is dim=-2 in ONNX.
-    dim = [d - 1 if d < 0 else d for d in dim]
-    return _fftn_onnx(self, dim, normalization, inverse=not forward, onesided=False)
+    assert(dim[-1] in dim == 2, "Unexpected input size")
+
+    signal = self
+    self_rank = len(self.shape)
+    signal_size = op.Size(signal)
+
+    # ONNX DFT input assumes the last dimension is the complex dimension.
+    # Thus dim=-1 in PyTorch is dim=-2 in ONNX.
+    dim = [(d - 1) + self_rank if d < 0 else d for d in dim]
+
+    transformed = signal
+
+    for dimension in reversed(dim):
+        transformed = op.DFT(transformed, axis=dimension, inverse=not forward, onesided=False)
+        if forward:
+            transformed = _fftn_onnx_normalization(transformed, normalization, signal_size)
+        else:
+            transformed = _fftn_onnx_inverse_normalization(transformed, normalization, signal_size)
+
+    # Unsure if output format is correct
+    return transformed
 
 
 @torch_op("aten::_fft_c2r", trace_only=True, complex=True)
 def aten__fft_c2r(
     self: TFloat,
     dim: Sequence[int],
     normalization: int,
-    last_dim_size: INT64,  # pylint: disable=unused-argument
+    last_dim_size: INT64,
 ) -> TFloat:
     """_fft_c2r(Tensor self, int[] dim, int normalization, SymInt last_dim_size) -> Tensor
 
     Complex to real inverse FFT.
     """
+    assert(dim[-1] in dim == 2, "Unexpected input size")
path: onnxscript/function_libs/torch_lib/ops/fft.py
user: login: github-advanced-security[bot]
body: ## RUFF/F631

Assert test is a non-empty tuple, which is always `True`.
See https://docs.astral.sh/ruff/rules/assert-tuple

[Show more details](https://github.com/microsoft/onnxscript/security/code-scanning/16521)
created_at: 2025-03-27T17:18:23Z
author_association: CONTRIBUTOR
diff_hunk: @@ -124,39 +67,74 @@
     Standard complex to complex FFT (forward or backward).
     """
 
-    # NOTE: trace_only because we need to negate forward
-    # NOTE: SymInt dim is not support because DFT-17 needs a static axis
-    # TODO(justinchuby): Make dim dynamic and remove trace_only when ONNX provides support
+    # NOTE: SymInt dim is not supported because DFT-17 needs a static axis
 
     # ONNX DFT input assumes the last dimension is the complex dimension.
     # Thus dim=-1 in PyTorch is dim=-2 in ONNX.
-    dim = [d - 1 if d < 0 else d for d in dim]
-    return _fftn_onnx(self, dim, normalization, inverse=not forward, onesided=False)
+    assert(dim[-1] in dim == 2, "Unexpected input size")
+
+    signal = self
+    self_rank = len(self.shape)
+    signal_size = op.Size(signal)
+
+    # ONNX DFT input assumes the last dimension is the complex dimension.
+    # Thus dim=-1 in PyTorch is dim=-2 in ONNX.
+    dim = [(d - 1) + self_rank if d < 0 else d for d in dim]
+
+    transformed = signal
+
+    for dimension in reversed(dim):
+        transformed = op.DFT(transformed, axis=dimension, inverse=not forward, onesided=False)
+        if forward:
+            transformed = _fftn_onnx_normalization(transformed, normalization, signal_size)
+        else:
+            transformed = _fftn_onnx_inverse_normalization(transformed, normalization, signal_size)
+
+    # Unsure if output format is correct
+    return transformed
 
 
 @torch_op("aten::_fft_c2r", trace_only=True, complex=True)
 def aten__fft_c2r(
     self: TFloat,
     dim: Sequence[int],
     normalization: int,
-    last_dim_size: INT64,  # pylint: disable=unused-argument
+    last_dim_size: INT64,
 ) -> TFloat:
     """_fft_c2r(Tensor self, int[] dim, int normalization, SymInt last_dim_size) -> Tensor
 
     Complex to real inverse FFT.
     """
+    assert(dim[-1] in dim == 2, "Unexpected input size")
path: onnxscript/function_libs/torch_lib/ops/fft.py
user: login: github-advanced-security[bot]
body: ## PYLINT/W3601

Suspicious 2-part chained comparison using semantically incompatible operators ('==' and 'in') (bad-chained-comparison)
See [bad-chained-comparison](https://pylint.pycqa.org/en/latest/user_guide/messages/warning/bad-chained-comparison.html).

To disable, use `  # pylint: disable=bad-chained-comparison`

[Show more details](https://github.com/microsoft/onnxscript/security/code-scanning/16525)
created_at: 2025-03-27T17:18:23Z
author_association: CONTRIBUTOR
diff_hunk: @@ -21,98 +21,44 @@
 from onnxscript.onnx_types import TensorType
 
 
-@torch_op(
-    ("aten::_fft_c2c", "aten::_fft_c2r", "aten::_fft_r2c"),
-    private=True,
-    complex=True,
-    trace_only=True,
-)
 def _fftn_onnx_normalization(
-    self,
-    transformed: TFloat,
+    self: TFloat,
     normalization: int,
-    forward: bool,
-    dims: Sequence[int],
+    signal_size: INT64,
 ) -> TFloat:
-    # Obtain the total_sample_count (n) for normalization
-    self_shape = op.Shape(self)
-    total_sample_count = op.ReduceProd(op.Gather(self_shape, dims), keepdims=0)
-    total_sample_count = op.CastLike(total_sample_count, transformed)
-
-    # Normalize the result
-    # Reference https://pytorch.org/docs/stable/generated/torch.fft.fftn.html#torch.fft.fftn
-    # Reference https://github.com/pytorch/pytorch/blob/d090c18fcaaba6e1b5cb474a89058cf6081c8275/torch/_refs/fft.py#L42
+    """ """
path: onnxscript/function_libs/torch_lib/ops/fft.py
user: login: github-advanced-security[bot]
body: ## RUFF/D419

Docstring is empty.
See https://docs.astral.sh/ruff/rules/empty-docstring

[Show more details](https://github.com/microsoft/onnxscript/security/code-scanning/16528)
created_at: 2025-03-27T17:29:21Z
author_association: CONTRIBUTOR
diff_hunk: @@ -21,98 +21,44 @@
 from onnxscript.onnx_types import TensorType
 
 
-@torch_op(
-    ("aten::_fft_c2c", "aten::_fft_c2r", "aten::_fft_r2c"),
-    private=True,
-    complex=True,
-    trace_only=True,
-)
 def _fftn_onnx_normalization(
-    self,
-    transformed: TFloat,
+    self: TFloat,
     normalization: int,
-    forward: bool,
-    dims: Sequence[int],
+    signal_size: INT64,
 ) -> TFloat:
-    # Obtain the total_sample_count (n) for normalization
-    self_shape = op.Shape(self)
-    total_sample_count = op.ReduceProd(op.Gather(self_shape, dims), keepdims=0)
-    total_sample_count = op.CastLike(total_sample_count, transformed)
-
-    # Normalize the result
-    # Reference https://pytorch.org/docs/stable/generated/torch.fft.fftn.html#torch.fft.fftn
-    # Reference https://github.com/pytorch/pytorch/blob/d090c18fcaaba6e1b5cb474a89058cf6081c8275/torch/_refs/fft.py#L42
+    """ """
+    # TODO: Make more efficient
+    # Norm values defined in https://github.com/pytorch/pytorch/blob/758d78790164bfb041555daed380de96e06f78a3/aten/src/ATen/native/SpectralOps.cpp#L117-L131
+    # Norm modes: https://github.com/pytorch/pytorch/blob/758d78790164bfb041555daed380de96e06f78a3/aten/src/ATen/native/SpectralOpsUtils.h#L15-L19
+    # Modes:
+    # 0: no normalization (backward)
+    # 1: "ortho" - divide by 1/sqrt(signal_size) (ortho)
+    # 2: divide by signal_size (forward)
     if normalization == 1:
-        # "forward" - normalize by 1/n
-        if forward:
-            result = op.Div(transformed, op.Sqrt(total_sample_count))
-        else:
-            result = op.Mul(transformed, op.Sqrt(total_sample_count))
+        self = op.Div(self, op.Sqrt(signal_size))
     elif normalization == 2:
-        # "ortho" - normalize by 1/sqrt(n)
-        if forward:
-            result = op.Div(transformed, total_sample_count)
-        else:
-            result = transformed
-    else:
-        # "backward" - no normalization
-        if forward:
-            result = transformed
-        else:
-            result = op.Mul(transformed, total_sample_count)
-
-    return result
+        self = op.Div(self, signal_size)
+    return self
 
 
-@torch_op(
-    ("aten::_fft_c2c", "aten::_fft_c2r", "aten::_fft_r2c"),
-    trace_only=True,
-    private=True,
-    complex=True,
-)
-def _fftn_onnx(
-    self: TFloat, dims: Sequence[int], normalization: int, inverse: bool, onesided: bool
+def _fftn_onnx_inverse_normalization(
+    self: TFloat,
+    normalization: int,
+    signal_size: INT64,
 ) -> TFloat:
-    """Standard complex to complex or real to complex FFT (forward or backward).
-
-    This is a private shared function for implementing the various FFT functions.
-
-    Args:
-        self: The input tensor.
-        dims: The dimensions to apply FFT.
-        normalization: The normalization mode.
-        inverse: Whether to compute the inverse FFT.
-        onesided: Whether to compute the one-sided FFT, which retains only the
-            positive frequencies.
-
-    Returns:
-        The transformed tensor.
-    """
-
-    # NOTE: trace_only because we need to process each dimension in a loop
-    # NOTE: SymInt dim is not support because DFT-17 needs a static axis
-    # TODO(justinchuby): Make dim dynamic and remove trace_only when ONNX provides support
-
-    # The 0-th dimension in ONNX DFT-17 is the batch dimension. We need to add a new
-    # dimension at the beginning to represent the batch dimension.
-    transformed = op.Unsqueeze(self, axes=[0])
-
-    # Add 1 to account for the batch dimension when counting axes from the left
-    new_dims = [dim_ + 1 if dim_ >= 0 else dim_ for dim_ in dims]
-
-    for dim in new_dims[:-1]:
-        transformed = op.DFT(transformed, axis=dim, inverse=inverse, onesided=False)
-
-    # Torch computers one-sided FFT on the last dimension only.
-    if onesided:
-        transformed = op.DFT(transformed, axis=new_dims[-1], inverse=inverse, onesided=True)
-    else:
-        transformed = op.DFT(transformed, axis=new_dims[-1], inverse=inverse, onesided=False)
-
-    # Remove the batch dimension
-    transformed = op.Squeeze(transformed, axes=[0])
-
-    return _fftn_onnx_normalization(self, transformed, normalization, not inverse, dims)
+    """ """
path: onnxscript/function_libs/torch_lib/ops/fft.py
user: login: github-advanced-security[bot]
body: ## RUFF/D419

Docstring is empty.
See https://docs.astral.sh/ruff/rules/empty-docstring

[Show more details](https://github.com/microsoft/onnxscript/security/code-scanning/16529)
created_at: 2025-03-27T17:29:21Z
author_association: CONTRIBUTOR
diff_hunk: @@ -124,39 +70,78 @@
     Standard complex to complex FFT (forward or backward).
     """
 
-    # NOTE: trace_only because we need to negate forward
-    # NOTE: SymInt dim is not support because DFT-17 needs a static axis
-    # TODO(justinchuby): Make dim dynamic and remove trace_only when ONNX provides support
+    # NOTE: SymInt dim is not supported because DFT-17 needs a static axis
+
+    # ONNX DFT input assumes the last dimension is the complex dimension.
+    # Thus dim=-1 in PyTorch is dim=-2 in ONNX.
+    assert (dim[-1] in dim == 2, "Unexpected input size")
path: onnxscript/function_libs/torch_lib/ops/fft.py
user: login: github-advanced-security[bot]
body: ## PYLINT/W3601

Suspicious 2-part chained comparison using semantically incompatible operators ('==' and 'in') (bad-chained-comparison)
See [bad-chained-comparison](https://pylint.pycqa.org/en/latest/user_guide/messages/warning/bad-chained-comparison.html).

To disable, use `  # pylint: disable=bad-chained-comparison`

[Show more details](https://github.com/microsoft/onnxscript/security/code-scanning/16526)
created_at: 2025-03-27T17:29:21Z
author_association: CONTRIBUTOR
diff_hunk: @@ -124,39 +70,78 @@
     Standard complex to complex FFT (forward or backward).
     """
 
-    # NOTE: trace_only because we need to negate forward
-    # NOTE: SymInt dim is not support because DFT-17 needs a static axis
-    # TODO(justinchuby): Make dim dynamic and remove trace_only when ONNX provides support
+    # NOTE: SymInt dim is not supported because DFT-17 needs a static axis
+
+    # ONNX DFT input assumes the last dimension is the complex dimension.
+    # Thus dim=-1 in PyTorch is dim=-2 in ONNX.
+    assert (dim[-1] in dim == 2, "Unexpected input size")
+
+    signal = self
+    self_rank = len(self.shape)
+    signal_size = op.Size(signal)
 
     # ONNX DFT input assumes the last dimension is the complex dimension.
     # Thus dim=-1 in PyTorch is dim=-2 in ONNX.
-    dim = [d - 1 if d < 0 else d for d in dim]
-    return _fftn_onnx(self, dim, normalization, inverse=not forward, onesided=False)
+    dim = [(d - 1) + self_rank if d < 0 else d for d in dim]
+
+    transformed = signal
+
+    for dimension in reversed(dim):
+        transformed = op.DFT(transformed, axis=dimension, inverse=not forward, onesided=False)
+        if forward:
+            transformed = _fftn_onnx_normalization(transformed, normalization, signal_size)
+        else:
+            transformed = _fftn_onnx_inverse_normalization(
+                transformed, normalization, signal_size
+            )
+
+    # Unsure if output format is correct
+    return transformed
 
 
 @torch_op("aten::_fft_c2r", trace_only=True, complex=True)
 def aten__fft_c2r(
     self: TFloat,
     dim: Sequence[int],
     normalization: int,
-    last_dim_size: INT64,  # pylint: disable=unused-argument
+    last_dim_size: INT64,
 ) -> TFloat:
     """_fft_c2r(Tensor self, int[] dim, int normalization, SymInt last_dim_size) -> Tensor
 
     Complex to real inverse FFT.
     """
+    assert (dim[-1] in dim == 2, "Unexpected input size")
path: onnxscript/function_libs/torch_lib/ops/fft.py
user: login: github-advanced-security[bot]
body: ## PYLINT/W3601

Suspicious 2-part chained comparison using semantically incompatible operators ('==' and 'in') (bad-chained-comparison)
See [bad-chained-comparison](https://pylint.pycqa.org/en/latest/user_guide/messages/warning/bad-chained-comparison.html).

To disable, use `  # pylint: disable=bad-chained-comparison`

[Show more details](https://github.com/microsoft/onnxscript/security/code-scanning/16527)
created_at: 2025-03-27T17:29:21Z
author_association: CONTRIBUTOR
diff_hunk: @@ -21,98 +21,44 @@
 from onnxscript.onnx_types import TensorType
 
 
-@torch_op(
-    ("aten::_fft_c2c", "aten::_fft_c2r", "aten::_fft_r2c"),
-    private=True,
-    complex=True,
-    trace_only=True,
-)
 def _fftn_onnx_normalization(
-    self,
-    transformed: TFloat,
+    self: TFloat,
     normalization: int,
-    forward: bool,
-    dims: Sequence[int],
+    signal_size: INT64,
path: onnxscript/function_libs/torch_lib/ops/fft.py
user: login: gramalingam
body: Question for @justinchuby : what does INT64 type signify? Specifically, is there a convention that tells us whether something is a statically known int value, or a symint, or a dynamic int value?
created_at: 2025-03-27T19:28:32Z
author_association: COLLABORATOR
diff_hunk: @@ -21,98 +21,44 @@
 from onnxscript.onnx_types import TensorType
 
 
-@torch_op(
-    ("aten::_fft_c2c", "aten::_fft_c2r", "aten::_fft_r2c"),
-    private=True,
-    complex=True,
-    trace_only=True,
-)
 def _fftn_onnx_normalization(
-    self,
-    transformed: TFloat,
+    self: TFloat,
     normalization: int,
-    forward: bool,
-    dims: Sequence[int],
+    signal_size: INT64,
path: onnxscript/function_libs/torch_lib/ops/fft.py
user: login: bmehta001
body: I saw that `op.Size` returns an int64, so I was using `INT64`. However, from the test errors, I realized that I need to cast `signal_size` to be the same type as `self` so that I can properly call op.Divide/op.Multiply with self/signal_size being the same type. 

I am not sure if this means `signal_size` should be `float` or `TFloat`.
created_at: 2025-03-27T20:29:34Z
author_association: CONTRIBUTOR
diff_hunk: @@ -21,98 +21,44 @@
 from onnxscript.onnx_types import TensorType
 
 
-@torch_op(
-    ("aten::_fft_c2c", "aten::_fft_c2r", "aten::_fft_r2c"),
-    private=True,
-    complex=True,
-    trace_only=True,
-)
 def _fftn_onnx_normalization(
-    self,
-    transformed: TFloat,
+    self: TFloat,
     normalization: int,
-    forward: bool,
-    dims: Sequence[int],
+    signal_size: INT64,
path: onnxscript/function_libs/torch_lib/ops/fft.py
user: login: titaiwangms
body: Is `signal_size` an output of an ONNX op? Output from an ONNX op should be Tensor, so it should be TFloat.
created_at: 2025-03-27T21:54:40Z
author_association: CONTRIBUTOR
diff_hunk: @@ -21,98 +21,44 @@
 from onnxscript.onnx_types import TensorType
 
 
-@torch_op(
-    ("aten::_fft_c2c", "aten::_fft_c2r", "aten::_fft_r2c"),
-    private=True,
-    complex=True,
-    trace_only=True,
-)
 def _fftn_onnx_normalization(
-    self,
-    transformed: TFloat,
+    self: TFloat,
     normalization: int,
-    forward: bool,
-    dims: Sequence[int],
+    signal_size: INT64,
path: onnxscript/function_libs/torch_lib/ops/fft.py
user: login: titaiwangms
body: We check https://github.com/pytorch/pytorch/blob/68414512e6fe641b02a6fe217fd516b7b776ea0d/aten/src/ATen/native/native_functions.yaml#L2991 for the signature of functions.
created_at: 2025-03-27T21:56:25Z
author_association: CONTRIBUTOR
diff_hunk: @@ -21,98 +21,44 @@
 from onnxscript.onnx_types import TensorType
 
 
-@torch_op(
-    ("aten::_fft_c2c", "aten::_fft_c2r", "aten::_fft_r2c"),
-    private=True,
-    complex=True,
-    trace_only=True,
-)
 def _fftn_onnx_normalization(
-    self,
-    transformed: TFloat,
+    self: TFloat,
     normalization: int,
-    forward: bool,
-    dims: Sequence[int],
+    signal_size: INT64,
path: onnxscript/function_libs/torch_lib/ops/fft.py
user: login: justinchuby
body: It’s usually symint. We use the same INT64 for symint and tensors
created_at: 2025-03-28T02:06:26Z
author_association: COLLABORATOR
diff_hunk: @@ -21,6 +21,12 @@
 from onnxscript.onnx_types import TensorType
 
 
+@torch_op(
+    ("aten::_fft_c2c", "aten::_fft_r2c"),
+    private=True,
path: onnxscript/function_libs/torch_lib/ops/fft.py
user: login: justinchuby
body: No need to add the torch_op decorator to private functions now. Thanks
created_at: 2025-03-31T16:24:42Z
author_association: COLLABORATOR
diff_hunk: @@ -21,6 +21,12 @@
 from onnxscript.onnx_types import TensorType
 
 
+@torch_op(
+    ("aten::_fft_c2c", "aten::_fft_r2c"),
+    private=True,
path: onnxscript/function_libs/torch_lib/ops/fft.py
user: login: titaiwangms
body: Sorry, that's me told him to do so. Sorry, @bmehta001 
created_at: 2025-03-31T16:28:40Z
author_association: CONTRIBUTOR
diff_hunk: @@ -684,25 +684,74 @@
 
 def sample_inputs__fft_c2r(self, device, dtype, requires_grad=False, **_):
     del self  # Unused
-    oned_tensor, nd_tensor = _prepare_data_for_fft_ops(device, dtype, requires_grad)
-
-    for normalization in (0, 1, 2):
+    if dtype == torch.complex128:
+        real_dtype = torch.float64
+    elif dtype == torch.complex64:
+        real_dtype = torch.float32
+    elif dtype == torch.complex32:
+        real_dtype = torch.float16
+    oned_tensor, nd_tensor = _prepare_data_for_fft_ops(device, real_dtype, requires_grad)
path: tests/function_libs/torch_lib/extra_opinfo.py
user: login: github-advanced-security[bot]
body: ## PYLINT/W0612

Unused variable 'nd_tensor' (unused-variable)
See [unused-variable](https://pylint.pycqa.org/en/latest/user_guide/messages/warning/unused-variable.html).

To disable, use `  # pylint: disable=unused-variable`

[Show more details](https://github.com/microsoft/onnxscript/security/code-scanning/16632)
created_at: 2025-04-10T17:28:32Z
author_association: CONTRIBUTOR
diff_hunk: @@ -684,25 +684,74 @@
 
 def sample_inputs__fft_c2r(self, device, dtype, requires_grad=False, **_):
     del self  # Unused
-    oned_tensor, nd_tensor = _prepare_data_for_fft_ops(device, dtype, requires_grad)
-
-    for normalization in (0, 1, 2):
+    if dtype == torch.complex128:
+        real_dtype = torch.float64
+    elif dtype == torch.complex64:
+        real_dtype = torch.float32
+    elif dtype == torch.complex32:
+        real_dtype = torch.float16
+    oned_tensor, nd_tensor = _prepare_data_for_fft_ops(device, real_dtype, requires_grad)
path: tests/function_libs/torch_lib/extra_opinfo.py
user: login: github-advanced-security[bot]
body: ## PYLINT/E0606

Possibly using variable 'real_dtype' before assignment (possibly-used-before-assignment)
See [possibly-used-before-assignment](https://pylint.pycqa.org/en/latest/user_guide/messages/error/possibly-used-before-assignment.html).

To disable, use `  # pylint: disable=possibly-used-before-assignment`

[Show more details](https://github.com/microsoft/onnxscript/security/code-scanning/16631)
created_at: 2025-04-10T17:28:33Z
author_association: CONTRIBUTOR
diff_hunk: @@ -684,25 +684,74 @@
 
 def sample_inputs__fft_c2r(self, device, dtype, requires_grad=False, **_):
     del self  # Unused
-    oned_tensor, nd_tensor = _prepare_data_for_fft_ops(device, dtype, requires_grad)
-
-    for normalization in (0, 1, 2):
+    if dtype == torch.complex128:
+        real_dtype = torch.float64
+    elif dtype == torch.complex64:
+        real_dtype = torch.float32
+    elif dtype == torch.complex32:
+        real_dtype = torch.float16
+    oned_tensor, nd_tensor = _prepare_data_for_fft_ops(device, real_dtype, requires_grad)
+
+    oned_tensor_result = oned_tensor()
+    complex_oned_tensor = functools.partial(
+        torch.fft.fft,
+        oned_tensor_result,
+        dim=0,
+    )
+    # nd_tensor_result = nd_tensor()
+    # for normalization in (0, 1, 2):
+    for normalization in (0,):
+        oned_tensor, nd_tensor = _prepare_data_for_fft_ops(device, real_dtype, requires_grad)
+        complex_oned_tensor = functools.partial(
+            torch.fft.fft,
+            oned_tensor().detach().clone(),
+            dim=0,
+        )
         # 1-D
         yield opinfo_core.SampleInput(
-            oned_tensor(), dim=(0,), normalization=normalization, last_dim_size=12
+            complex_oned_tensor(),
+            dim=(0,),
+            normalization=normalization,
+            last_dim_size=31,
         )
-        # N-D
-        for dim in [
-            (0,),
-            (1,),
-            (2,),
-            (1, 2),
-            (0, 1),
-            (0, 1, 2),
-        ]:
-            yield opinfo_core.SampleInput(
-                nd_tensor(), dim=dim, normalization=normalization, last_dim_size=6
-            )
+        # One-sided: Refers to the one-sided FFT result, which includes only the positive frequency terms.
+        # # One-sided
+        oned_tensor, nd_tensor = _prepare_data_for_fft_ops(device, real_dtype, requires_grad)
path: tests/function_libs/torch_lib/extra_opinfo.py
user: login: github-advanced-security[bot]
body: ## RUFF/RUF059

Unpacked variable `nd_tensor` is never used.
See https://docs.astral.sh/ruff/rules/unused-unpacked-variable

[Show more details](https://github.com/microsoft/onnxscript/security/code-scanning/16633)
created_at: 2025-04-10T17:28:33Z
author_association: CONTRIBUTOR
diff_hunk: @@ -684,24 +684,34 @@
 
 def sample_inputs__fft_c2r(self, device, dtype, requires_grad=False, **_):
     del self  # Unused
-    oned_tensor, nd_tensor = _prepare_data_for_fft_ops(device, dtype, requires_grad)
-
+    real_dtype = torch.float
+    if dtype == torch.cdouble:
+        real_dtype = torch.double
+    oned_tensor, nd_tensor = _prepare_data_for_fft_ops(device, real_dtype, requires_grad)
+    oned_tensor_result = oned_tensor()
+    nd_tensor_result = nd_tensor()
+    complex_oned_tensor = torch.ops.aten._fft_r2c.default(oned_tensor_result, [0], normalization=0, onesided=False)
path: tests/function_libs/torch_lib/extra_opinfo.py
user: login: github-advanced-security[bot]
body: ## PYLINT/W0212

Access to a protected member _fft_r2c of a client class (protected-access)
See [protected-access](https://pylint.pycqa.org/en/latest/user_guide/messages/warning/protected-access.html).

To disable, use `  # pylint: disable=protected-access`

[Show more details](https://github.com/microsoft/onnxscript/security/code-scanning/16635)
created_at: 2025-04-11T01:11:01Z
author_association: CONTRIBUTOR
diff_hunk: @@ -684,24 +684,34 @@
 
 def sample_inputs__fft_c2r(self, device, dtype, requires_grad=False, **_):
     del self  # Unused
-    oned_tensor, nd_tensor = _prepare_data_for_fft_ops(device, dtype, requires_grad)
-
+    real_dtype = torch.float
+    if dtype == torch.cdouble:
+        real_dtype = torch.double
+    oned_tensor, nd_tensor = _prepare_data_for_fft_ops(device, real_dtype, requires_grad)
+    oned_tensor_result = oned_tensor()
+    nd_tensor_result = nd_tensor()
+    complex_oned_tensor = torch.ops.aten._fft_r2c.default(oned_tensor_result, [0], normalization=0, onesided=False)
+    # for normalization in (0, 1, 2):
     for normalization in (0, 1, 2):
         # 1-D
         yield opinfo_core.SampleInput(
-            oned_tensor(), dim=(0,), normalization=normalization, last_dim_size=12
+            complex_oned_tensor,
+            dim=(0,),
+            normalization=normalization,
+            last_dim_size=31,
         )
         # N-D
         for dim in [
             (0,),
             (1,),
             (2,),
-            (1, 2),
-            (0, 1),
-            (0, 1, 2),
         ]:
+            complex_nd_tensor = torch.ops.aten._fft_r2c.default(nd_tensor_result, dim, normalization=0, onesided=False)
path: tests/function_libs/torch_lib/extra_opinfo.py
user: login: github-advanced-security[bot]
body: ## PYLINT/W0212

Access to a protected member _fft_r2c of a client class (protected-access)
See [protected-access](https://pylint.pycqa.org/en/latest/user_guide/messages/warning/protected-access.html).

To disable, use `  # pylint: disable=protected-access`

[Show more details](https://github.com/microsoft/onnxscript/security/code-scanning/16636)
created_at: 2025-04-11T01:11:01Z
author_association: CONTRIBUTOR
diff_hunk: @@ -124,39 +70,95 @@ def aten__fft_c2c(
     Standard complex to complex FFT (forward or backward).
     """
 
-    # NOTE: trace_only because we need to negate forward
-    # NOTE: SymInt dim is not support because DFT-17 needs a static axis
-    # TODO(justinchuby): Make dim dynamic and remove trace_only when ONNX provides support
+    # NOTE: SymInt dim is not supported because DFT-17 needs a static axis
+
+    # ONNX DFT input assumes the last dimension is the complex dimension.
+    # Thus dim=-1 in PyTorch is dim=-2 in ONNX.
+    self_rank = len(self.shape)
 
     # ONNX DFT input assumes the last dimension is the complex dimension.
     # Thus dim=-1 in PyTorch is dim=-2 in ONNX.
-    dim = [d - 1 if d < 0 else d for d in dim]
-    return _fftn_onnx(self, dim, normalization, inverse=not forward, onesided=False)
+    dim = [(d - 1) + self_rank if d < 0 else d for d in dim]
+
+    unsqueeze_first_dim = 0 in dim
path: onnxscript/function_libs/torch_lib/ops/fft.py
user: login: justinchuby
body: nice in saving the unsqueeze for common cases!
created_at: 2025-04-11T01:29:04Z
author_association: COLLABORATOR
diff_hunk: @@ -684,24 +684,34 @@ def sample_inputs__fft_r2c(self, device, dtype, requires_grad=False, **_):
 
 def sample_inputs__fft_c2r(self, device, dtype, requires_grad=False, **_):
     del self  # Unused
-    oned_tensor, nd_tensor = _prepare_data_for_fft_ops(device, dtype, requires_grad)
-
+    real_dtype = torch.float
+    if dtype == torch.cdouble:
path: tests/function_libs/torch_lib/extra_opinfo.py
user: login: justinchuby
body: ```suggestion
    if dtype == torch.complex128:
```
created_at: 2025-04-11T01:31:44Z
author_association: COLLABORATOR
diff_hunk: @@ -124,39 +70,103 @@ def aten__fft_c2c(
     Standard complex to complex FFT (forward or backward).
     """
 
-    # NOTE: trace_only because we need to negate forward
-    # NOTE: SymInt dim is not support because DFT-17 needs a static axis
-    # TODO(justinchuby): Make dim dynamic and remove trace_only when ONNX provides support
+    # NOTE: SymInt dim is not supported because DFT-17 needs a static axis
 
     # ONNX DFT input assumes the last dimension is the complex dimension.
     # Thus dim=-1 in PyTorch is dim=-2 in ONNX.
-    dim = [d - 1 if d < 0 else d for d in dim]
-    return _fftn_onnx(self, dim, normalization, inverse=not forward, onesided=False)
+    self_rank = len(self.shape)
+
+    # ONNX DFT input assumes the last dimension is the complex dimension.
+    # Thus dim=-1 in PyTorch is dim=-2 in ONNX.
+    dim = [(d - 1) + self_rank if d < 0 else d for d in dim]
+
+    unsqueeze_first_dim = 0 in dim
+    if unsqueeze_first_dim:
+        transformed = op.Unsqueeze(self, axes=[0])
+        # Add 1 to account for the batch dimension when counting axes from the left
+        dim = [dim_ + 1 if dim_ >= 0 else dim_ for dim_ in dim]
+    else:
+        transformed = self
+
+    for dimension in reversed(dim):
+        transformed = op.DFT(transformed, axis=dimension, inverse=not forward, onesided=False)
+        if forward:
+            transformed = _fftn_onnx_normalization(
+                transformed,
+                normalization,
+                op.CastLike(self.shape[dimension - unsqueeze_first_dim], transformed),
+            )
+        else:
+            transformed = _fftn_onnx_inverse_normalization(
+                transformed,
+                normalization,
+                op.CastLike(self.shape[dimension - unsqueeze_first_dim], transformed),
path: onnxscript/function_libs/torch_lib/ops/fft.py
user: login: justinchuby
body: op.shape
created_at: 2025-04-11T15:12:01Z
author_association: COLLABORATOR
diff_hunk: @@ -124,39 +70,103 @@ def aten__fft_c2c(
     Standard complex to complex FFT (forward or backward).
     """
 
-    # NOTE: trace_only because we need to negate forward
-    # NOTE: SymInt dim is not support because DFT-17 needs a static axis
-    # TODO(justinchuby): Make dim dynamic and remove trace_only when ONNX provides support
+    # NOTE: SymInt dim is not supported because DFT-17 needs a static axis
 
     # ONNX DFT input assumes the last dimension is the complex dimension.
     # Thus dim=-1 in PyTorch is dim=-2 in ONNX.
-    dim = [d - 1 if d < 0 else d for d in dim]
-    return _fftn_onnx(self, dim, normalization, inverse=not forward, onesided=False)
+    self_rank = len(self.shape)
+
+    # ONNX DFT input assumes the last dimension is the complex dimension.
+    # Thus dim=-1 in PyTorch is dim=-2 in ONNX.
+    dim = [(d - 1) + self_rank if d < 0 else d for d in dim]
+
+    unsqueeze_first_dim = 0 in dim
+    if unsqueeze_first_dim:
+        transformed = op.Unsqueeze(self, axes=[0])
+        # Add 1 to account for the batch dimension when counting axes from the left
+        dim = [dim_ + 1 if dim_ >= 0 else dim_ for dim_ in dim]
+    else:
+        transformed = self
+
+    for dimension in reversed(dim):
+        transformed = op.DFT(transformed, axis=dimension, inverse=not forward, onesided=False)
+        if forward:
+            transformed = _fftn_onnx_normalization(
+                transformed,
+                normalization,
+                op.CastLike(self.shape[dimension - unsqueeze_first_dim], transformed),
+            )
+        else:
+            transformed = _fftn_onnx_inverse_normalization(
+                transformed,
+                normalization,
+                op.CastLike(self.shape[dimension - unsqueeze_first_dim], transformed),
+            )
+
+    if unsqueeze_first_dim:
+        transformed = op.Squeeze(transformed, axes=[0])
+
+    return transformed
 
 
 @torch_op("aten::_fft_c2r", trace_only=True, complex=True)
 def aten__fft_c2r(
     self: TFloat,
     dim: Sequence[int],
     normalization: int,
-    last_dim_size: INT64,  # pylint: disable=unused-argument
+    last_dim_size: INT64,
 ) -> TFloat:
     """_fft_c2r(Tensor self, int[] dim, int normalization, SymInt last_dim_size) -> Tensor
 
     Complex to real inverse FFT.
     """
-
-    # TODO(justinchuby): Figure out what last_dim_size does
+    if len(dim) != 1:
+        raise NotImplementedError("Only one dimension is supported for inverse FFT")
 
     self_rank = len(self.shape)
+
     # ONNX DFT input assumes the last dimension is the complex dimension.
     # Thus dim=-1 in PyTorch is dim=-2 in ONNX.
     dim = [(d - 1) + self_rank if d < 0 else d for d in dim]
-    transformed = _fftn_onnx(self, dim, normalization, inverse=True, onesided=False)
-    # Take only the real part
-    real_part = op.Slice(transformed, axes=[-1], starts=[0], ends=[1])
 
-    return op.Squeeze(real_part, axes=[-1])
+    unsqueeze_first_dim = 0 in dim
+    if unsqueeze_first_dim:
+        transformed = op.Unsqueeze(self, axes=[0])
+        # Add 1 to account for the batch dimension when counting axes from the left
+        dim = [dim_ + 1 if dim_ >= 0 else dim_ for dim_ in dim]
+    else:
+        transformed = self
+
+    for idx, dimension in enumerate(reversed(dim)):
+        if idx > 0:
+            transformed = op.DFT(transformed, axis=dimension, inverse=True, onesided=False)
+        else:
+            # Torch truncates/pads on the last dimension only. Typically, the only valid values that can be passed
+            # into PyTorch are n or n//2+1, where n is self.shape[dim[-1]], but this is not always the case, so we
+            # place no such restriction on the ONNX side.
+            transformed = op.DFT(
+                transformed,
+                dft_length=last_dim_size,
+                axis=dimension,
+                inverse=True,
+                onesided=False,
+            )
+        transformed = _fftn_onnx_inverse_normalization(
+            transformed,
+            normalization,
+            op.CastLike(
path: onnxscript/function_libs/torch_lib/ops/fft.py
user: login: justinchuby
body: maybe move castlike into _fftn_onnx_inverse_normalization? I feel that can make this more readable.
created_at: 2025-04-11T15:13:19Z
author_association: COLLABORATOR
diff_hunk: @@ -21,98 +21,33 @@
 from onnxscript.onnx_types import TensorType
 
 
-@torch_op(
-    ("aten::_fft_c2c", "aten::_fft_c2r", "aten::_fft_r2c"),
-    private=True,
-    complex=True,
-    trace_only=True,
-)
 def _fftn_onnx_normalization(
-    self,
-    transformed: TFloat,
+    self: TFloat,
     normalization: int,
-    forward: bool,
-    dims: Sequence[int],
-) -> TFloat:
-    # Obtain the total_sample_count (n) for normalization
-    self_shape = op.Shape(self)
-    total_sample_count = op.ReduceProd(op.Gather(self_shape, dims), keepdims=0)
-    total_sample_count = op.CastLike(total_sample_count, transformed)
-
-    # Normalize the result
-    # Reference https://pytorch.org/docs/stable/generated/torch.fft.fftn.html#torch.fft.fftn
-    # Reference https://github.com/pytorch/pytorch/blob/d090c18fcaaba6e1b5cb474a89058cf6081c8275/torch/_refs/fft.py#L42
-    if normalization == 1:
-        # "forward" - normalize by 1/n
-        if forward:
-            result = op.Div(transformed, op.Sqrt(total_sample_count))
-        else:
-            result = op.Mul(transformed, op.Sqrt(total_sample_count))
-    elif normalization == 2:
-        # "ortho" - normalize by 1/sqrt(n)
-        if forward:
-            result = op.Div(transformed, total_sample_count)
-        else:
-            result = transformed
-    else:
-        # "backward" - no normalization
-        if forward:
-            result = transformed
-        else:
-            result = op.Mul(transformed, total_sample_count)
-
-    return result
-
-
-@torch_op(
-    ("aten::_fft_c2c", "aten::_fft_c2r", "aten::_fft_r2c"),
-    trace_only=True,
-    private=True,
-    complex=True,
-)
-def _fftn_onnx(
-    self: TFloat, dims: Sequence[int], normalization: int, inverse: bool, onesided: bool
+    signal_size,
path: onnxscript/function_libs/torch_lib/ops/fft.py
user: login: justinchuby
body: ```suggestion
    signal_size: INT64,
```

nit
created_at: 2025-04-13T14:43:36Z
author_association: COLLABORATOR
diff_hunk: @@ -21,98 +21,33 @@
 from onnxscript.onnx_types import TensorType
 
 
-@torch_op(
-    ("aten::_fft_c2c", "aten::_fft_c2r", "aten::_fft_r2c"),
-    private=True,
-    complex=True,
-    trace_only=True,
-)
 def _fftn_onnx_normalization(
-    self,
-    transformed: TFloat,
+    self: TFloat,
     normalization: int,
-    forward: bool,
-    dims: Sequence[int],
-) -> TFloat:
-    # Obtain the total_sample_count (n) for normalization
-    self_shape = op.Shape(self)
-    total_sample_count = op.ReduceProd(op.Gather(self_shape, dims), keepdims=0)
-    total_sample_count = op.CastLike(total_sample_count, transformed)
-
-    # Normalize the result
-    # Reference https://pytorch.org/docs/stable/generated/torch.fft.fftn.html#torch.fft.fftn
-    # Reference https://github.com/pytorch/pytorch/blob/d090c18fcaaba6e1b5cb474a89058cf6081c8275/torch/_refs/fft.py#L42
-    if normalization == 1:
-        # "forward" - normalize by 1/n
-        if forward:
-            result = op.Div(transformed, op.Sqrt(total_sample_count))
-        else:
-            result = op.Mul(transformed, op.Sqrt(total_sample_count))
-    elif normalization == 2:
-        # "ortho" - normalize by 1/sqrt(n)
-        if forward:
-            result = op.Div(transformed, total_sample_count)
-        else:
-            result = transformed
-    else:
-        # "backward" - no normalization
-        if forward:
-            result = transformed
-        else:
-            result = op.Mul(transformed, total_sample_count)
-
-    return result
-
-
-@torch_op(
-    ("aten::_fft_c2c", "aten::_fft_c2r", "aten::_fft_r2c"),
-    trace_only=True,
-    private=True,
-    complex=True,
-)
-def _fftn_onnx(
-    self: TFloat, dims: Sequence[int], normalization: int, inverse: bool, onesided: bool
+    signal_size,
path: onnxscript/function_libs/torch_lib/ops/fft.py
user: login: bmehta001
body: Fixed
created_at: 2025-04-14T08:27:37Z
author_association: CONTRIBUTOR
diff_hunk: @@ -124,39 +70,103 @@ def aten__fft_c2c(
     Standard complex to complex FFT (forward or backward).
     """
 
-    # NOTE: trace_only because we need to negate forward
-    # NOTE: SymInt dim is not support because DFT-17 needs a static axis
-    # TODO(justinchuby): Make dim dynamic and remove trace_only when ONNX provides support
+    # NOTE: SymInt dim is not supported because DFT-17 needs a static axis
 
     # ONNX DFT input assumes the last dimension is the complex dimension.
     # Thus dim=-1 in PyTorch is dim=-2 in ONNX.
-    dim = [d - 1 if d < 0 else d for d in dim]
-    return _fftn_onnx(self, dim, normalization, inverse=not forward, onesided=False)
+    self_rank = len(self.shape)
+
+    # ONNX DFT input assumes the last dimension is the complex dimension.
+    # Thus dim=-1 in PyTorch is dim=-2 in ONNX.
+    dim = [(d - 1) + self_rank if d < 0 else d for d in dim]
+
+    unsqueeze_first_dim = 0 in dim
+    if unsqueeze_first_dim:
+        transformed = op.Unsqueeze(self, axes=[0])
+        # Add 1 to account for the batch dimension when counting axes from the left
+        dim = [dim_ + 1 if dim_ >= 0 else dim_ for dim_ in dim]
+    else:
+        transformed = self
+
+    for dimension in reversed(dim):
+        transformed = op.DFT(transformed, axis=dimension, inverse=not forward, onesided=False)
+        if forward:
+            transformed = _fftn_onnx_normalization(
+                transformed,
+                normalization,
+                op.CastLike(self.shape[dimension - unsqueeze_first_dim], transformed),
+            )
+        else:
+            transformed = _fftn_onnx_inverse_normalization(
+                transformed,
+                normalization,
+                op.CastLike(self.shape[dimension - unsqueeze_first_dim], transformed),
+            )
+
+    if unsqueeze_first_dim:
+        transformed = op.Squeeze(transformed, axes=[0])
+
+    return transformed
 
 
 @torch_op("aten::_fft_c2r", trace_only=True, complex=True)
 def aten__fft_c2r(
     self: TFloat,
     dim: Sequence[int],
     normalization: int,
-    last_dim_size: INT64,  # pylint: disable=unused-argument
+    last_dim_size: INT64,
 ) -> TFloat:
     """_fft_c2r(Tensor self, int[] dim, int normalization, SymInt last_dim_size) -> Tensor
 
     Complex to real inverse FFT.
     """
-
-    # TODO(justinchuby): Figure out what last_dim_size does
+    if len(dim) != 1:
+        raise NotImplementedError("Only one dimension is supported for inverse FFT")
 
     self_rank = len(self.shape)
+
     # ONNX DFT input assumes the last dimension is the complex dimension.
     # Thus dim=-1 in PyTorch is dim=-2 in ONNX.
     dim = [(d - 1) + self_rank if d < 0 else d for d in dim]
-    transformed = _fftn_onnx(self, dim, normalization, inverse=True, onesided=False)
-    # Take only the real part
-    real_part = op.Slice(transformed, axes=[-1], starts=[0], ends=[1])
 
-    return op.Squeeze(real_part, axes=[-1])
+    unsqueeze_first_dim = 0 in dim
+    if unsqueeze_first_dim:
+        transformed = op.Unsqueeze(self, axes=[0])
+        # Add 1 to account for the batch dimension when counting axes from the left
+        dim = [dim_ + 1 if dim_ >= 0 else dim_ for dim_ in dim]
+    else:
+        transformed = self
+
+    for idx, dimension in enumerate(reversed(dim)):
+        if idx > 0:
+            transformed = op.DFT(transformed, axis=dimension, inverse=True, onesided=False)
+        else:
+            # Torch truncates/pads on the last dimension only. Typically, the only valid values that can be passed
+            # into PyTorch are n or n//2+1, where n is self.shape[dim[-1]], but this is not always the case, so we
+            # place no such restriction on the ONNX side.
+            transformed = op.DFT(
+                transformed,
+                dft_length=last_dim_size,
+                axis=dimension,
+                inverse=True,
+                onesided=False,
+            )
+        transformed = _fftn_onnx_inverse_normalization(
+            transformed,
+            normalization,
+            op.CastLike(
path: onnxscript/function_libs/torch_lib/ops/fft.py
user: login: bmehta001
body: Done + merged normalized functions to simplify a bit
created_at: 2025-04-14T08:28:23Z
author_association: CONTRIBUTOR
diff_hunk: @@ -124,39 +70,103 @@ def aten__fft_c2c(
     Standard complex to complex FFT (forward or backward).
     """
 
-    # NOTE: trace_only because we need to negate forward
-    # NOTE: SymInt dim is not support because DFT-17 needs a static axis
-    # TODO(justinchuby): Make dim dynamic and remove trace_only when ONNX provides support
+    # NOTE: SymInt dim is not supported because DFT-17 needs a static axis
 
     # ONNX DFT input assumes the last dimension is the complex dimension.
     # Thus dim=-1 in PyTorch is dim=-2 in ONNX.
-    dim = [d - 1 if d < 0 else d for d in dim]
-    return _fftn_onnx(self, dim, normalization, inverse=not forward, onesided=False)
+    self_rank = len(self.shape)
+
+    # ONNX DFT input assumes the last dimension is the complex dimension.
+    # Thus dim=-1 in PyTorch is dim=-2 in ONNX.
+    dim = [(d - 1) + self_rank if d < 0 else d for d in dim]
+
+    unsqueeze_first_dim = 0 in dim
+    if unsqueeze_first_dim:
+        transformed = op.Unsqueeze(self, axes=[0])
+        # Add 1 to account for the batch dimension when counting axes from the left
+        dim = [dim_ + 1 if dim_ >= 0 else dim_ for dim_ in dim]
+    else:
+        transformed = self
+
+    for dimension in reversed(dim):
+        transformed = op.DFT(transformed, axis=dimension, inverse=not forward, onesided=False)
+        if forward:
+            transformed = _fftn_onnx_normalization(
+                transformed,
+                normalization,
+                op.CastLike(self.shape[dimension - unsqueeze_first_dim], transformed),
+            )
+        else:
+            transformed = _fftn_onnx_inverse_normalization(
+                transformed,
+                normalization,
+                op.CastLike(self.shape[dimension - unsqueeze_first_dim], transformed),
path: onnxscript/function_libs/torch_lib/ops/fft.py
user: login: bmehta001
body: Fixed everywhere
created_at: 2025-04-14T08:28:32Z
author_association: CONTRIBUTOR
diff_hunk: @@ -124,39 +70,95 @@ def aten__fft_c2c(
     Standard complex to complex FFT (forward or backward).
     """
 
-    # NOTE: trace_only because we need to negate forward
-    # NOTE: SymInt dim is not support because DFT-17 needs a static axis
-    # TODO(justinchuby): Make dim dynamic and remove trace_only when ONNX provides support
+    # NOTE: SymInt dim is not supported because DFT-17 needs a static axis
+
+    # ONNX DFT input assumes the last dimension is the complex dimension.
+    # Thus dim=-1 in PyTorch is dim=-2 in ONNX.
+    self_rank = len(self.shape)
 
     # ONNX DFT input assumes the last dimension is the complex dimension.
     # Thus dim=-1 in PyTorch is dim=-2 in ONNX.
-    dim = [d - 1 if d < 0 else d for d in dim]
-    return _fftn_onnx(self, dim, normalization, inverse=not forward, onesided=False)
+    dim = [(d - 1) + self_rank if d < 0 else d for d in dim]
+
+    unsqueeze_first_dim = 0 in dim
path: onnxscript/function_libs/torch_lib/ops/fft.py
user: login: bmehta001
body: Thanks, this was something I took from your PR 😅 I should have started working off of your branch in the first place - I tried to rebase off of it now, but there were too many merge conflicts to make the commit history smooth.

I did change something that I originally borrowed from your PR - accounting for negative dimensions. The original PyTorch function does not support negative dimensions, so I dropped support for them (unless it's better to support additional cases beyond what the original torch ops function supports).
created_at: 2025-04-14T08:31:08Z
author_association: CONTRIBUTOR
diff_hunk: @@ -124,39 +70,95 @@ def aten__fft_c2c(
     Standard complex to complex FFT (forward or backward).
     """
 
-    # NOTE: trace_only because we need to negate forward
-    # NOTE: SymInt dim is not support because DFT-17 needs a static axis
-    # TODO(justinchuby): Make dim dynamic and remove trace_only when ONNX provides support
+    # NOTE: SymInt dim is not supported because DFT-17 needs a static axis
+
+    # ONNX DFT input assumes the last dimension is the complex dimension.
+    # Thus dim=-1 in PyTorch is dim=-2 in ONNX.
+    self_rank = len(self.shape)
 
     # ONNX DFT input assumes the last dimension is the complex dimension.
     # Thus dim=-1 in PyTorch is dim=-2 in ONNX.
-    dim = [d - 1 if d < 0 else d for d in dim]
-    return _fftn_onnx(self, dim, normalization, inverse=not forward, onesided=False)
+    dim = [(d - 1) + self_rank if d < 0 else d for d in dim]
+
+    unsqueeze_first_dim = 0 in dim
path: onnxscript/function_libs/torch_lib/ops/fft.py
user: login: justinchuby
body: Ah so I was patting myself on the back 😅. Either way great work for making it work!
created_at: 2025-04-14T14:56:52Z
author_association: COLLABORATOR
diff_hunk: @@ -684,24 +684,38 @@ def sample_inputs__fft_r2c(self, device, dtype, requires_grad=False, **_):
 
 def sample_inputs__fft_c2r(self, device, dtype, requires_grad=False, **_):
     del self  # Unused
-    oned_tensor, nd_tensor = _prepare_data_for_fft_ops(device, dtype, requires_grad)
-
+    real_dtype = torch.float
+    if dtype == torch.complex128:
+        real_dtype = torch.double
+    oned_tensor, nd_tensor = _prepare_data_for_fft_ops(device, real_dtype, requires_grad)
+    oned_tensor_result = oned_tensor()
+    nd_tensor_result = nd_tensor()
+    complex_oned_tensor = torch.ops.aten._fft_r2c.default(  # pylint: disable=protected-access
+        oned_tensor_result, [0], normalization=0, onesided=False
+    )
+    # for normalization in (0, 1, 2):
     for normalization in (0, 1, 2):
         # 1-D
         yield opinfo_core.SampleInput(
-            oned_tensor(), dim=(0,), normalization=normalization, last_dim_size=12
+            complex_oned_tensor,
+            dim=(0,),
+            normalization=normalization,
+            last_dim_size=31,
         )
         # N-D
         for dim in [
             (0,),
             (1,),
             (2,),
-            (1, 2),
-            (0, 1),
-            (0, 1, 2),
path: tests/function_libs/torch_lib/extra_opinfo.py
user: login: justinchuby
body: Note to self: this should be added back in a follow up
created_at: 2025-04-16T01:34:26Z
author_association: COLLABORATOR

PR comments:
user: login: codecov[bot]
created_at: 2025-03-27T17:02:37Z
body: ## [Codecov](https://app.codecov.io/gh/microsoft/onnxscript/pull/2141?dropdown=coverage&src=pr&el=h1&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=microsoft) Report
Attention: Patch coverage is `96.00000%` with `2 lines` in your changes missing coverage. Please review.
> Project coverage is 74.27%. Comparing base [(`e404922`)](https://app.codecov.io/gh/microsoft/onnxscript/commit/e404922d0e8bacc0aa4a008f396efe4178f005fc?dropdown=coverage&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=microsoft) to head [(`f2555ca`)](https://app.codecov.io/gh/microsoft/onnxscript/commit/f2555cab03a87e96876cf2fa7a256d7f3d155efc?dropdown=coverage&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=microsoft).
> Report is 9 commits behind head on main.

| [Files with missing lines](https://app.codecov.io/gh/microsoft/onnxscript/pull/2141?dropdown=coverage&src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=microsoft) | Patch % | Lines |
|---|---|---|
| [onnxscript/function\_libs/torch\_lib/ops/fft.py](https://app.codecov.io/gh/microsoft/onnxscript/pull/2141?src=pr&el=tree&filepath=onnxscript%2Ffunction_libs%2Ftorch_lib%2Fops%2Ffft.py&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=microsoft#diff-b25ueHNjcmlwdC9mdW5jdGlvbl9saWJzL3RvcmNoX2xpYi9vcHMvZmZ0LnB5) | 96.00% | [1 Missing and 1 partial :warning: ](https://app.codecov.io/gh/microsoft/onnxscript/pull/2141?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=microsoft) |

<details><summary>Additional details and impacted files</summary>


```diff
@@            Coverage Diff             @@
##             main    #2141      +/-   ##
==========================================
+ Coverage   74.25%   74.27%   +0.02%     
==========================================
  Files         226      226              
  Lines       29442    29454      +12     
  Branches     3424     3432       +8     
==========================================
+ Hits        21862    21878      +16     
+ Misses       6424     6419       -5     
- Partials     1156     1157       +1     
```

</details>

[:umbrella: View full report in Codecov by Sentry](https://app.codecov.io/gh/microsoft/onnxscript/pull/2141?dropdown=coverage&src=pr&el=continue&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=microsoft).   
:loudspeaker: Have feedback on the report? [Share it here](https://about.codecov.io/codecov-pr-comment-feedback/?utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=microsoft).

author_association: NONE
user: login: titaiwangms
created_at: 2025-03-27T17:50:04Z
body: The error looks like a mismatch of complex (torch) vs real representation (ONNX). Maybe explore https://github.com/microsoft/onnxscript/blob/a36ec86d09350a14856364618f5888bb5073ccad/tests/function_libs/torch_lib/ops_test.py#L197

Or using https://github.com/microsoft/onnxscript/blob/a36ec86d09350a14856364618f5888bb5073ccad/tests/function_libs/torch_lib/ops_test.py#L22
author_association: CONTRIBUTOR
user: login: justinchuby
created_at: 2025-03-28T15:55:59Z
body: > Output type needs to be converted to Complex Tensor

output doesn't need to be converted. We always use the real repr for complex values in onnx. The torch exporter will keep track of this information.
author_association: COLLABORATOR
user: login: justinchuby
created_at: 2025-04-07T12:25:55Z
body: For the errors `RuntimeError: out_size == signal_size[i + 1] || out_size == (signal_size[i + 1] / 2) + 1 INTERNAL ASSERT FAILED at "../aten/src/ATen/native/mkl/SpectralOps.cpp":468, please report a bug to PyTorch` it's possible to change the test cases because current ones are invalid for pytorch.
author_association: COLLABORATOR
user: login: justinchuby
created_at: 2025-04-14T14:58:34Z
body: @bmehta001 not sure why the cla bot is complaining even though you have join the org. Can you comment following the second (company) option?
author_association: COLLABORATOR
user: login: justinchuby
created_at: 2025-04-14T16:28:21Z
body: Could you test (https://github.com/pytorch/pytorch/issues/119360)

```py
import torch

from torch import nn


class iRFFTModel(nn.Module):
  def __init__(self):
      super().__init__()

  def forward(self, x):
    return torch.fft.irfft(torch.view_as_complex(x), n=960)

example_tensor = torch.rand(481, 2)
irfft_model = iRFFTModel()
irfft_model_onnx_path = "irfft_model.onnx"
out_torch = irfft_model(example_tensor)
print("torch shape:", out_torch.shape)

torch.onnx.export(
    irfft_model,
    (example_tensor,),
    dynamo=True,
    verify=True
)
```

(https://github.com/pytorch/pytorch/issues/139246)

```py
import torch
from torch import nn
import torchaudio

class DataCov(nn.Module):
    def __init__(self):
        super(DataCov, self).__init__()

        self.transform = nn.Sequential(
            torchaudio.transforms.MelSpectrogram(
                sample_rate=48000, n_fft=1536, hop_length=768, f_min=20, f_max=20000
            )
        )

    def forward(self, x1):
        return self.transform(x1)


def load_data_cov():
    module = DataCov().to(torch.float32).to("cpu")
    module.eval()
    return module


data_cov = load_data_cov()
x = torch.randn((1, 1, 12 * 48000), dtype=torch.float32, device="cpu")
y = data_cov(x)
input_names = ["x"]
output_names = ["output"]

torch.onnx.export(
    data_cov,
    (x,),
    dynamo=True,
    verify=True,
)
```


(https://github.com/pytorch/pytorch/issues/113444)

```py
import torch
import torch.nn as nn


def fftconv(u, k, D):
    seqlen = u.shape[-1]
    fft_size = 2 * seqlen

    k_f = torch.fft.rfft(k, n=fft_size) / fft_size
    u_f = torch.fft.rfft(u.to(dtype=k.dtype), n=fft_size)

    if len(u.shape) > 3: k_f = k_f.unsqueeze(1)
    y = torch.fft.irfft(u_f * k_f, n=fft_size, norm='forward')[..., :seqlen]

    out = y + u * D.unsqueeze(-1)
    return out.to(dtype=u.dtype)


class Filter(nn.Module):
    def forward(self, x, k=None, bias=None):
        y = fftconv(x, k, bias)
        return y



filter = Filter().eval()
x_input, k_input, bias_input = torch.rand(1, 512, 1024), torch.rand(512, 1024), torch.rand(512)
export_output = torch.onnx.export(filter, (x_input, k_input, bias_input), dynamo=True, verify=True)
```
author_association: COLLABORATOR
user: login: bmehta001
created_at: 2025-04-14T17:02:13Z
body: > @bmehta001 please read the following Contributor License Agreement(CLA). If you agree with the CLA, please reply with the following information.
> 
> ```
> @microsoft-github-policy-service agree [company="{your company}"]
> ```
> 
> > Options:
> > 
> > * (default - no company specified) I have sole ownership of intellectual property rights to my Submissions and I am not making Submissions in the course of work for my employer.
> > 
> > ```
> > @microsoft-github-policy-service agree
> > ```
> > 
> > 
> >     
> >       
> >     
> > 
> >       
> >     
> > 
> >     
> >   
> > 
> > * (when company given) I am making Submissions in the course of work for my employer (or my employer has intellectual property rights in my Submissions by contract or applicable law). I have permission from my employer to make Submissions and enter into this Agreement on behalf of my employer. By signing below, the defined term “You” includes me and my employer.
> > 
> > ```
> > @microsoft-github-policy-service agree company="Microsoft"
> > ```
> 
> Contributor License Agreement
> ### Contribution License Agreement
> This Contribution License Agreement (**“Agreement”**) is agreed to by the party signing below (**“You”**), and conveys certain license rights to Microsoft Corporation and its affiliates (“Microsoft”) for Your contributions to Microsoft open source projects. This Agreement is effective as of the latest signature date below.
> 
> 1. **Definitions**.
>    **“Code”** means the computer software code, whether in human-readable or machine-executable form,
>    that is delivered by You to Microsoft under this Agreement.
>    **“Project”** means any of the projects owned or managed by Microsoft and offered under a license
>    approved by the Open Source Initiative ([www.opensource.org](http://www.opensource.org)).
>    **“Submit”** is the act of uploading, submitting, transmitting, or distributing code or other content to any
>    Project, including but not limited to communication on electronic mailing lists, source code control
>    systems, and issue tracking systems that are managed by, or on behalf of, the Project for the purpose of
>    discussing and improving that Project, but excluding communication that is conspicuously marked or
>    otherwise designated in writing by You as “Not a Submission.”
>    **“Submission”** means the Code and any other copyrightable material Submitted by You, including any
>    associated comments and documentation.
> 2. **Your Submission**. You must agree to the terms of this Agreement before making a Submission to any
>    Project. This Agreement covers any and all Submissions that You, now or in the future (except as
>    described in Section 4 below), Submit to any Project.
> 3. **Originality of Work**. You represent that each of Your Submissions is entirely Your original work.
>    Should You wish to Submit materials that are not Your original work, You may Submit them separately
>    to the Project if You (a) retain all copyright and license information that was in the materials as You
>    received them, (b) in the description accompanying Your Submission, include the phrase “Submission
>    containing materials of a third party:” followed by the names of the third party and any licenses or other
>    restrictions of which You are aware, and (c) follow any other instructions in the Project’s written
>    guidelines concerning Submissions.
> 4. **Your Employer**. References to “employer” in this Agreement include Your employer or anyone else
>    for whom You are acting in making Your Submission, e.g. as a contractor, vendor, or agent. If Your
>    Submission is made in the course of Your work for an employer or Your employer has intellectual
>    property rights in Your Submission by contract or applicable law, You must secure permission from Your
>    employer to make the Submission before signing this Agreement. In that case, the term “You” in this
>    Agreement will refer to You and the employer collectively. If You change employers in the future and
>    desire to Submit additional Submissions for the new employer, then You agree to sign a new Agreement
>    and secure permission from the new employer before Submitting those Submissions.
> 5. **Licenses**.
> 
> * **Copyright License**. You grant Microsoft, and those who receive the Submission directly or
>   indirectly from Microsoft, a perpetual, worldwide, non-exclusive, royalty-free, irrevocable license in the
>   Submission to reproduce, prepare derivative works of, publicly display, publicly perform, and distribute
>   the Submission and such derivative works, and to sublicense any or all of the foregoing rights to third
>   parties.
> * **Patent License**. You grant Microsoft, and those who receive the Submission directly or
>   indirectly from Microsoft, a perpetual, worldwide, non-exclusive, royalty-free, irrevocable license under
>   Your patent claims that are necessarily infringed by the Submission or the combination of the
>   Submission with the Project to which it was Submitted to make, have made, use, offer to sell, sell and
>   import or otherwise dispose of the Submission alone or with the Project.
> * **Other Rights Reserved**. Each party reserves all rights not expressly granted in this Agreement.
>   No additional licenses or rights whatsoever (including, without limitation, any implied licenses) are
>   granted by implication, exhaustion, estoppel or otherwise.
> 
> 6. **Representations and Warranties**. You represent that You are legally entitled to grant the above
>    licenses. You represent that each of Your Submissions is entirely Your original work (except as You may
>    have disclosed under Section 3). You represent that You have secured permission from Your employer to
>    make the Submission in cases where Your Submission is made in the course of Your work for Your
>    employer or Your employer has intellectual property rights in Your Submission by contract or applicable
>    law. If You are signing this Agreement on behalf of Your employer, You represent and warrant that You
>    have the necessary authority to bind the listed employer to the obligations contained in this Agreement.
>    You are not expected to provide support for Your Submission, unless You choose to do so. UNLESS
>    REQUIRED BY APPLICABLE LAW OR AGREED TO IN WRITING, AND EXCEPT FOR THE WARRANTIES
>    EXPRESSLY STATED IN SECTIONS 3, 4, AND 6, THE SUBMISSION PROVIDED UNDER THIS AGREEMENT IS
>    PROVIDED WITHOUT WARRANTY OF ANY KIND, INCLUDING, BUT NOT LIMITED TO, ANY WARRANTY OF
>    NONINFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE.
> 7. **Notice to Microsoft**. You agree to notify Microsoft in writing of any facts or circumstances of which
>    You later become aware that would make Your representations in this Agreement inaccurate in any
>    respect.
> 8. **Information about Submissions**. You agree that contributions to Projects and information about
>    contributions may be maintained indefinitely and disclosed publicly, including Your name and other
>    information that You submit with Your Submission.
> 9. **Governing Law/Jurisdiction**. This Agreement is governed by the laws of the State of Washington, and
>    the parties consent to exclusive jurisdiction and venue in the federal courts sitting in King County,
>    Washington, unless no federal subject matter jurisdiction exists, in which case the parties consent to
>    exclusive jurisdiction and venue in the Superior Court of King County, Washington. The parties waive all
>    defenses of lack of personal jurisdiction and forum non-conveniens.
> 10. **Entire Agreement/Assignment**. This Agreement is the entire agreement between the parties, and
>     supersedes any and all prior agreements, understandings or communications, written or oral, between
>     the parties relating to the subject matter hereof. This Agreement may be assigned by Microsoft.

@microsoft-github-policy-service agree company="Microsoft"
author_association: CONTRIBUTOR
user: login: justinchuby
created_at: 2025-04-16T01:33:10Z
body: I am going to take the liberty to merge this PR so that the main improvements can be made available. Please feel free to create follow ups when you get a chance to look at the pytorch examples. Thanks for the great work!
author_association: COLLABORATOR