Repository: pytorch/TensorRT (language: python)

Issue:
❓ [Question] operator being decomposed rather than being converted when a corresponding converter exists?
## ❓ Question

From the debug log below, it seems that the `aten.grid_sampler_2d` operator gets decomposed into several lower-level operators. But isn't there a corresponding [converter](https://github.com/pytorch/TensorRT/blob/9a100b6414bee175040bcaa275ecb71df54836e4/py/torch_tensorrt/dynamo/conversion/aten_ops_converters.py#L333-L358) which should be used?

## What you have already tried

```py
import torch
import torch.nn as nn
import torch.nn.functional as F
import torch_tensorrt


class MyModule(nn.Module):
    def __init__(self):
        super().__init__()
        
    def forward(self, input, grid):
        return F.grid_sample(input, grid, mode="bilinear", padding_mode="border", align_corners=True)
    
model = MyModule().eval().cuda()

inputs = [
    torch.randn((1, 3, 8, 8), dtype=torch.float, device="cuda"),
    torch.randn((1, 16, 16, 2), dtype=torch.float, device="cuda")
]

optimized_model = torch_tensorrt.compile(
    model,
    ir="dynamo",
    inputs=inputs,
    enabled_precisions={torch.float},
    debug=True,
    min_block_size=1,
    truncate_long_and_double=True,
    output_format="fx",
)
```

```
DEBUG:torch_tensorrt.dynamo.conversion.aten_ops_converters:_to_copy converter rejected node _to_copy with dtype torch.int64
DEBUG:torch_tensorrt.dynamo.conversion.aten_ops_converters:_to_copy converter rejected node _to_copy with dtype torch.int64
DEBUG:torch_tensorrt.dynamo.conversion.aten_ops_converters:_to_copy converter rejected node _to_copy_1 with dtype torch.int64
DEBUG:torch_tensorrt.dynamo.conversion.aten_ops_converters:_to_copy converter rejected node _to_copy_1 with dtype torch.int64
DEBUG:torch_tensorrt.dynamo.conversion.aten_ops_converters:_to_copy converter rejected node _to_copy_2 with dtype torch.int64
DEBUG:torch_tensorrt.dynamo.conversion.aten_ops_converters:_to_copy converter rejected node _to_copy_2 with dtype torch.int64
DEBUG:torch_tensorrt.dynamo.conversion.aten_ops_converters:_to_copy converter rejected node _to_copy_3 with dtype torch.int64
DEBUG:torch_tensorrt.dynamo.conversion.aten_ops_converters:_to_copy converter rejected node _to_copy_3 with dtype torch.int64
DEBUG:torch_tensorrt.dynamo.conversion.aten_ops_converters:_to_copy converter rejected node _to_copy_4 with dtype torch.int64
DEBUG:torch_tensorrt.dynamo.conversion.aten_ops_converters:_to_copy converter rejected node _to_copy_4 with dtype torch.int64
DEBUG:torch_tensorrt.dynamo.conversion.aten_ops_converters:_to_copy converter rejected node _to_copy_5 with dtype torch.int64
DEBUG:torch_tensorrt.dynamo.conversion.aten_ops_converters:_to_copy converter rejected node _to_copy_5 with dtype torch.int64
DEBUG:torch_tensorrt.dynamo.conversion.aten_ops_converters:_to_copy converter rejected node _to_copy_6 with dtype torch.int64
DEBUG:torch_tensorrt.dynamo.conversion.aten_ops_converters:_to_copy converter rejected node _to_copy_6 with dtype torch.int64
DEBUG:torch_tensorrt.dynamo.conversion.aten_ops_converters:_to_copy converter rejected node _to_copy_7 with dtype torch.int64
DEBUG:torch_tensorrt.dynamo.conversion.aten_ops_converters:_to_copy converter rejected node _to_copy_7 with dtype torch.int64
DEBUG:torch_tensorrt.dynamo.partitioning._global_partitioner:
Supported Nodes:
- torch.ops.aten.reshape.default + Operator Count: 13
- torch.ops.aten.expand.default + Operator Count: 1
- torch.ops.aten.select.int + Operator Count: 2
- torch.ops.aten.mul.Tensor + Operator Count: 10
- torch.ops.aten.add.Tensor + Operator Count: 7
- torch.ops.aten.clamp.default + Operator Count: 2
- torch.ops.aten.floor.default + Operator Count: 2
- torch.ops.aten.sub.Tensor + Operator Count: 8
- torch.ops.aten.ge.Scalar + Operator Count: 8
- torch.ops.aten.lt.Scalar + Operator Count: 8
- torch.ops.aten.logical_and.default + Operator Count: 12
- torch.ops.aten.where.self + Operator Count: 12
- torch.ops.aten.index.Tensor + Operator Count: 4

DEBUG:torch_tensorrt.dynamo.partitioning._global_partitioner:
Unsupported or Excluded Nodes:
- torch.ops.aten._to_copy.default + Operator Count: 8

DEBUG:torch_tensorrt.dynamo._compiler:Detected support for 89 operators out of 97 in subgraph.
DEBUG:torch_tensorrt.dynamo.conversion.aten_ops_converters:_to_copy converter rejected node _to_copy with dtype torch.int64
DEBUG:torch_tensorrt.dynamo.conversion.aten_ops_converters:_to_copy converter rejected node _to_copy with dtype torch.int64
DEBUG:torch_tensorrt.dynamo.conversion.aten_ops_converters:_to_copy converter rejected node _to_copy_1 with dtype torch.int64
DEBUG:torch_tensorrt.dynamo.conversion.aten_ops_converters:_to_copy converter rejected node _to_copy_1 with dtype torch.int64
DEBUG:torch_tensorrt.dynamo.conversion.aten_ops_converters:_to_copy converter rejected node _to_copy_2 with dtype torch.int64
DEBUG:torch_tensorrt.dynamo.conversion.aten_ops_converters:_to_copy converter rejected node _to_copy_2 with dtype torch.int64
DEBUG:torch_tensorrt.dynamo.conversion.aten_ops_converters:_to_copy converter rejected node _to_copy_3 with dtype torch.int64
DEBUG:torch_tensorrt.dynamo.conversion.aten_ops_converters:_to_copy converter rejected node _to_copy_3 with dtype torch.int64
DEBUG:torch_tensorrt.dynamo.conversion.aten_ops_converters:_to_copy converter rejected node _to_copy_4 with dtype torch.int64
DEBUG:torch_tensorrt.dynamo.conversion.aten_ops_converters:_to_copy converter rejected node _to_copy_4 with dtype torch.int64
DEBUG:torch_tensorrt.dynamo.conversion.aten_ops_converters:_to_copy converter rejected node _to_copy_5 with dtype torch.int64
DEBUG:torch_tensorrt.dynamo.conversion.aten_ops_converters:_to_copy converter rejected node _to_copy_5 with dtype torch.int64
DEBUG:torch_tensorrt.dynamo.conversion.aten_ops_converters:_to_copy converter rejected node _to_copy_6 with dtype torch.int64
DEBUG:torch_tensorrt.dynamo.conversion.aten_ops_converters:_to_copy converter rejected node _to_copy_6 with dtype torch.int64
DEBUG:torch_tensorrt.dynamo.conversion.aten_ops_converters:_to_copy converter rejected node _to_copy_7 with dtype torch.int64
DEBUG:torch_tensorrt.dynamo.conversion.aten_ops_converters:_to_copy converter rejected node _to_copy_7 with dtype torch.int64
DEBUG:torch_tensorrt.dynamo.partitioning._adjacency_partitioner:
Number of TensorRT-Accelerated Engines Generated: 2
DEBUG:torch_tensorrt.dynamo.partitioning._adjacency_partitioner:
Supported Nodes:
- torch.ops.aten.reshape.default + Operator Count: 13
- torch.ops.aten.expand.default + Operator Count: 1
- torch.ops.aten.select.int + Operator Count: 2
- torch.ops.aten.mul.Tensor + Operator Count: 10
- torch.ops.aten.add.Tensor + Operator Count: 7
- torch.ops.aten.clamp.default + Operator Count: 2
- torch.ops.aten.floor.default + Operator Count: 2
- torch.ops.aten.sub.Tensor + Operator Count: 8
- torch.ops.aten.ge.Scalar + Operator Count: 8
- torch.ops.aten.lt.Scalar + Operator Count: 8
- torch.ops.aten.logical_and.default + Operator Count: 12
- torch.ops.aten.where.self + Operator Count: 12
- torch.ops.aten.index.Tensor + Operator Count: 4

DEBUG:torch_tensorrt.dynamo.partitioning._adjacency_partitioner:
Unsupported or Excluded Nodes:
- torch.ops.aten._to_copy.default + Operator Count: 8

++++++++++++++++++++++++++++++++++++++++++++++++++ Dry-Run Results for Graph ++++++++++++++++++++++++++++++++++++++++++++++++++

The graph consists of 97 Total Operators, of which 89 operators are supported, 91.75% coverage

The following ops are currently unsupported or excluded from conversion, and are listed with their op-count in the graph:
 torch.ops.aten._to_copy.default: 8

The following nodes are currently set to run in Torch:
Node: torch.ops.aten._to_copy.default, with layer location: __/_to_copy
Node: torch.ops.aten._to_copy.default, with layer location: __/_to_copy_1
Node: torch.ops.aten._to_copy.default, with layer location: __/_to_copy_2
Node: torch.ops.aten._to_copy.default, with layer location: __/_to_copy_3
Node: torch.ops.aten._to_copy.default, with layer location: __/_to_copy_4
Node: torch.ops.aten._to_copy.default, with layer location: __/_to_copy_5
Node: torch.ops.aten._to_copy.default, with layer location: __/_to_copy_6
Node: torch.ops.aten._to_copy.default, with layer location: __/_to_copy_7
Note: Some of the above nodes may be supported, but were not included in a TRT graph by the partitioner

Compiled with: CompilationSettings(precision=torch.float32, debug=True, workspace_size=0, min_block_size=1, torch_executed_ops=set(), pass_through_build_failures=False, max_aux_streams=None, version_compatible=False, optimization_level=None, use_python_runtime=False, truncate_long_and_double=True, use_fast_partitioner=True, enable_experimental_decompositions=False, device=Device(type=DeviceType.GPU, gpu_id=0), require_full_compilation=False, disable_tf32=False, sparse_weights=False, refit=False, engine_capability=<EngineCapability.DEFAULT: 0>, num_avg_timing_iters=1, dla_sram_size=1048576, dla_local_dram_size=1073741824, dla_global_dram_size=536870912, dryrun=False, hardware_compatible=False, output_format='fx')

  Graph Structure:

   Inputs: List[Tensor: (1, 3, 8, 8)@float32, Tensor: (1, 16, 16, 2)@float32]
    ...
    TRT Engine #1 - Submodule name: _run_on_acc_0
     Engine Inputs: List[Tensor: (1, 16, 16, 2)@float32]
     Number of Operators in Engine: 58
     Engine Outputs: Tuple(Tensor: (1, 3, 16, 16)@float32, Tensor: (1, 3, 16, 16)@float32, Tensor: (1, 3, 16, 16)@bool, Tensor: (1, 3, 16, 16)@float32, Tensor: (1, 3, 16, 16)@float32, Tensor: (1, 3, 16, 16)@bool, Tensor: (1, 3, 16, 16)@float32, Tensor: (1, 3, 16, 16)@float32, Tensor: (1, 3, 16, 16)@bool, Tensor: (1, 3, 16, 16)@float32, Tensor: (1, 3, 16, 16)@bool, Tensor: (1, 3, 16, 16)@float32)
    ...
    TRT Engine #2 - Submodule name: _run_on_acc_2
     Engine Inputs: List[Tensor: (1, 3, 16, 16)@bool, Tensor: (1, 3, 16, 16)@int32, Tensor: (1, 3, 16, 16)@int32, Tensor: (1, 3, 8, 8)@float32, Tensor: (1, 3, 16, 16)@float32, Tensor: (1, 3, 16, 16)@bool, Tensor: (1, 3, 16, 16)@int32, Tensor: (1, 3, 16, 16)@int32, Tensor: (1, 3, 16, 16)@float32, Tensor: (1, 3, 16, 16)@bool, Tensor: (1, 3, 16, 16)@int32, Tensor: (1, 3, 16, 16)@int32, Tensor: (1, 3, 16, 16)@float32, Tensor: (1, 3, 16, 16)@bool, Tensor: (1, 3, 16, 16)@int32, Tensor: (1, 3, 16, 16)@int32, Tensor: (1, 3, 16, 16)@float32]
     Number of Operators in Engine: 31
     Engine Outputs: Tensor: (1, 3, 16, 16)@float32
    ...
   Outputs: List[Tensor: (1, 3, 16, 16)@float32]

  ------------------------- Aggregate Stats -------------------------

   Average Number of Operators per TRT Engine: 44.5
   Most Operators in a TRT Engine: 58

  ********** Recommendations **********

   - For minimal graph segmentation, select min_block_size=58 which would generate 1 TRT engine(s)
   - For moderate graph segmentation, select min_block_size=45 which would generate 1 TRT engine(s)
   - The current level of graph segmentation is equivalent to selecting min_block_size=31 which generates 2 TRT engine(s)
```

## Environment

 - PyTorch Version (e.g., 1.0): 2.3.0.dev20240221+cu121
 - CPU Architecture: x64
 - OS (e.g., Linux): Ubuntu 22.04 LTS
 - How you installed PyTorch (`conda`, `pip`, `libtorch`, source): pip
 - Build command you used (if compiling from source):
 - Are you using local sources or building from archives:
 - Python version: 3.10.12
 - CUDA version: 12.1
 - GPU models and configuration: RTX 3050
 - Any other relevant information:


Discussion / hints:
The same for [aten.leaky_relu](https://github.com/pytorch/TensorRT/blob/9a100b6414bee175040bcaa275ecb71df54836e4/py/torch_tensorrt/dynamo/conversion/aten_ops_converters.py#L447-L462).

```py
import torch
import torch.nn as nn
import torch_tensorrt

class MyModule(nn.Module):
    def __init__(self):
        super().__init__()
        self.m = nn.LeakyReLU()

    def forward(self, x):
        return self.m(x)

model = MyModule().eval().cuda().half()
inputs = [torch.randn((1, 3, 4, 4), dtype=torch.half, device="cuda")]

optimized_model = torch_tensorrt.compile(
    model,
    ir="dynamo",
    inputs=inputs,
    enabled_precisions={torch.half},
    debug=True,
    min_block_size=1,
)
```

```
DEBUG:torch_tensorrt.dynamo.partitioning._global_partitioner:
Supported Nodes:
- torch.ops.aten._to_copy.default + Operator Count: 2
- torch.ops.aten.gt.Scalar + Operator Count: 1
- torch.ops.aten.mul.Tensor + Operator Count: 1
- torch.ops.aten.where.self + Operator Count: 1

DEBUG:torch_tensorrt.dynamo.partitioning._global_partitioner:
All Nodes Supported

DEBUG:torch_tensorrt.dynamo._compiler:Detected support for 5 operators out of 5 in subgraph.
DEBUG:torch_tensorrt.dynamo.partitioning._adjacency_partitioner:
Number of TensorRT-Accelerated Engines Generated: 1
DEBUG:torch_tensorrt.dynamo.partitioning._adjacency_partitioner:
Supported Nodes:
- torch.ops.aten._to_copy.default + Operator Count: 2
- torch.ops.aten.gt.Scalar + Operator Count: 1
- torch.ops.aten.mul.Tensor + Operator Count: 1
- torch.ops.aten.where.self + Operator Count: 1

DEBUG:torch_tensorrt.dynamo.partitioning._adjacency_partitioner:
All Nodes Supported

++++++++++++++++++++++++++++++++++++++++++++++++++ Dry-Run Results for Graph ++++++++++++++++++++++++++++++++++++++++++++++++++

The graph consists of 5 Total Operators, of which 5 operators are supported, 100.0% coverage

Compiled with: CompilationSettings(precision=torch.float16, debug=True, workspace_size=0, min_block_size=1, torch_executed_ops=set(), pass_through_build_failures=False, max_aux_streams=None, version_compatible=False, optimization_level=None, use_python_runtime=False, truncate_long_and_double=False, use_fast_partitioner=True, enable_experimental_decompositions=False, device=Device(type=DeviceType.GPU, gpu_id=0), require_full_compilation=False, disable_tf32=False, sparse_weights=False, refit=False, engine_capability=<EngineCapability.DEFAULT: 0>, num_avg_timing_iters=1, dla_sram_size=1048576, dla_local_dram_size=1073741824, dla_global_dram_size=536870912, dryrun=False, hardware_compatible=False, output_format='exported_program')

  Graph Structure:

   Inputs: List[Tensor: (1, 3, 4, 4)@float16]
    ...
    TRT Engine #1 - Submodule name: _run_on_acc_0
     Engine Inputs: List[Tensor: (1, 3, 4, 4)@float16]
     Number of Operators in Engine: 5
     Engine Outputs: Tensor: (1, 3, 4, 4)@float16
    ...
   Outputs: List[Tensor: (1, 3, 4, 4)@float16]

  ------------------------- Aggregate Stats -------------------------

   Average Number of Operators per TRT Engine: 5.0
   Most Operators in a TRT Engine: 5

  ********** Recommendations **********

   - For minimal graph segmentation, select min_block_size=5 which would generate 1 TRT engine(s)
   - The current level of graph segmentation is equivalent to selecting min_block_size=5 which generates 1 TRT engine(s)
WARNING: [Torch-TensorRT] - Using default stream in enqueue()/enqueueV2()/enqueueV3() may lead to performance issues due to additional cudaDeviceSynchronize() calls by TensorRT to ensure correct synchronizations. Please use non-default stream instead.
C:\Python311\Lib\site-packages\torch\export\exported_program.py:740: UserWarning: Unable to execute the generated python source code from the graph. The graph module will no longer be directly callable, but you can still run the ExportedProgram, and if needed, you can run the graph module eagerly using torch.fx.Interpreter.
  warnings.warn(
```
[aten.upsample_bilinear2d](https://github.com/pytorch/TensorRT/blob/main/py/torch_tensorrt/dynamo/conversion/aten_ops_converters.py#L2515-L2533)

```py
import torch
import torch.nn as nn
import torch.nn.functional as F
import torch_tensorrt


class MyModule(nn.Module):
    def __init__(self):
        super().__init__()

    def forward(self, x):
        return F.interpolate(x, scale_factor=2, mode="bilinear", align_corners=True)


model = MyModule().eval().cuda().half()
inputs = [
    torch.randn((1, 3, 128, 128), dtype=torch.half, device="cuda"),
]

optimized_model = torch_tensorrt.compile(
    model,
    ir="dynamo",
    inputs=inputs,
    enabled_precisions={torch.half},
    debug=True,
    min_block_size=1,
)
```

```
DEBUG:torch_tensorrt.dynamo.partitioning._global_partitioner:
Supported Nodes:
- torch.ops.aten._to_copy.default + Operator Count: 2
- torch.ops.aten.index.Tensor + Operator Count: 4
- torch.ops.aten.sub.Tensor + Operator Count: 3
- torch.ops.aten.mul.Tensor + Operator Count: 3
- torch.ops.aten.add.Tensor + Operator Count: 3

DEBUG:torch_tensorrt.dynamo.partitioning._global_partitioner:
All Nodes Supported

DEBUG:torch_tensorrt.dynamo._compiler:Detected support for 15 operators out of 15 in subgraph.
DEBUG:torch_tensorrt.dynamo.partitioning._adjacency_partitioner:
Number of TensorRT-Accelerated Engines Generated: 1
DEBUG:torch_tensorrt.dynamo.partitioning._adjacency_partitioner:
Supported Nodes:
- torch.ops.aten._to_copy.default + Operator Count: 2
- torch.ops.aten.index.Tensor + Operator Count: 4
- torch.ops.aten.sub.Tensor + Operator Count: 3
- torch.ops.aten.mul.Tensor + Operator Count: 3
- torch.ops.aten.add.Tensor + Operator Count: 3

DEBUG:torch_tensorrt.dynamo.partitioning._adjacency_partitioner:
All Nodes Supported

++++++++++++++++++++++++++++++++++++++++++++++++++ Dry-Run Results for Graph ++++++++++++++++++++++++++++++++++++++++++++++++++

The graph consists of 15 Total Operators, of which 15 operators are supported, 100.0% coverage

Compiled with: CompilationSettings(precision=torch.float16, debug=True, workspace_size=0, min_block_size=1, torch_executed_ops=set(), pass_through_build_failures=False, max_aux_streams=None, version_compatible=False, optimization_level=None, use_python_runtime=False, truncate_long_and_double=True, use_fast_partitioner=True, enable_experimental_decompositions=False, device=Device(type=DeviceType.GPU, gpu_id=0), require_full_compilation=False, disable_tf32=False, sparse_weights=False, refit=False, engine_capability=<EngineCapability.DEFAULT: 0>, num_avg_timing_iters=1, dla_sram_size=1048576, dla_local_dram_size=1073741824, dla_global_dram_size=536870912, dryrun=False, hardware_compatible=False, output_format='exported_program')

  Graph Structure:

   Inputs: List[Tensor: (1, 3, 128, 128)@float16]
    ...
    TRT Engine #1 - Submodule name: _run_on_acc_0
     Engine Inputs: List[Tensor: (1, 3, 128, 128)@float16]
     Number of Operators in Engine: 15
     Engine Outputs: Tensor: (1, 3, 256, 256)@float16
    ...
   Outputs: List[Tensor: (1, 3, 256, 256)@float16]

  ------------------------- Aggregate Stats -------------------------

   Average Number of Operators per TRT Engine: 15.0
   Most Operators in a TRT Engine: 15

  ********** Recommendations **********

   - For minimal graph segmentation, select min_block_size=15 which would generate 1 TRT engine(s)
   - The current level of graph segmentation is equivalent to selecting min_block_size=15 which generates 1 TRT engine(s)
WARNING: [Torch-TensorRT] - Using default stream in enqueue()/enqueueV2()/enqueueV3() may lead to performance issues due to additional cudaDeviceSynchronize() calls by TensorRT to ensure correct synchronizations. Please use non-default stream instead.
C:\Python311\Lib\site-packages\torch\export\exported_program.py:740: UserWarning: Unable to execute the generated python source code from the graph. The graph module will no longer be directly callable, but you can still run the ExportedProgram, and if needed, you can run the graph module eagerly using torch.fx.Interpreter.
  warnings.warn(
```

Fix patch:
diff --git a/py/torch_tensorrt/dynamo/conversion/aten_ops_converters.py b/py/torch_tensorrt/dynamo/conversion/aten_ops_converters.py
index 9acb750aed..0ea2f0304d 100644
--- a/py/torch_tensorrt/dynamo/conversion/aten_ops_converters.py
+++ b/py/torch_tensorrt/dynamo/conversion/aten_ops_converters.py
@@ -332,6 +332,8 @@ def aten_ops_fmod(
 
 @dynamo_tensorrt_converter(torch.ops.aten.grid_sampler)
 @dynamo_tensorrt_converter(torch.ops.aten.grid_sampler_2d)
+@dynamo_tensorrt_converter(torch.ops.aten.grid_sampler.default)
+@dynamo_tensorrt_converter(torch.ops.aten.grid_sampler_2d.default)
 @enforce_tensor_types(
     {
         0: (TRTTensor,),
diff --git a/py/torch_tensorrt/dynamo/lowering/_decomposition_groups.py b/py/torch_tensorrt/dynamo/lowering/_decomposition_groups.py
index de791851db..d365d6d9c5 100644
--- a/py/torch_tensorrt/dynamo/lowering/_decomposition_groups.py
+++ b/py/torch_tensorrt/dynamo/lowering/_decomposition_groups.py
@@ -46,7 +46,6 @@
     aten.gelu,
     aten.gelu_backward,
     aten.glu_backward,
-    aten.grid_sampler_2d,
     aten.hardshrink,
     aten.hardshrink_backward,
     aten.hardsigmoid,


Test patch:
diff --git a/tests/py/dynamo/conversion/test_grid_aten.py b/tests/py/dynamo/conversion/test_grid_aten.py
index 32480110f3..e3b5783b19 100644
--- a/tests/py/dynamo/conversion/test_grid_aten.py
+++ b/tests/py/dynamo/conversion/test_grid_aten.py
@@ -6,112 +6,74 @@
 from torch.testing._internal.common_utils import run_tests
 from torch_tensorrt import Input
 
+grid_sampler_aten_ops = {
+    "torch.ops.aten.grid_sampler": torch.ops.aten.grid_sampler,
+    "torch.ops.aten.grid_sampler_2d": torch.ops.aten.grid_sampler_2d,
+    "torch.ops.aten.grid_sampler.default": torch.ops.aten.grid_sampler.default,
+    "torch.ops.aten.grid_sampler_2d.default": torch.ops.aten.grid_sampler_2d.default,
+}
+
 grid_sampler_ops = [
     (
         "input_grid_interpolation_nearest_sample_fill",
-        (lambda x, grid: torch.ops.aten.grid_sampler(x, grid, 0, 0, True)),
+        "torch.ops.aten.grid_sampler",
+        (lambda x, grid, op: op(x, grid, 0, 0, True)),
         [1, 1, 5, 5],
         [1, 5, 2, 2],
     ),
     (
         "input_grid_interpolation_nearest_sample_clamp",
-        (lambda x, grid: torch.ops.aten.grid_sampler(x, grid, 0, 1, True)),
+        "torch.ops.aten.grid_sampler",
+        (lambda x, grid, op: op(x, grid, 0, 1, True)),
         [1, 1, 5, 5],
         [1, 5, 2, 2],
     ),
     (
         "input_grid_interpolation_nearest_sample_reflect",
-        (lambda x, grid: torch.ops.aten.grid_sampler(x, grid, 0, 2, True)),
+        "torch.ops.aten.grid_sampler",
+        (lambda x, grid, op: op(x, grid, 0, 2, True)),
         [1, 1, 5, 5],
         [1, 5, 2, 2],
     ),
     (
         "input_grid_interpolation_linear_sample_fill",
-        (lambda x, grid: torch.ops.aten.grid_sampler(x, grid, 0, 0, True)),
+        "torch.ops.aten.grid_sampler",
+        (lambda x, grid, op: op(x, grid, 1, 0, True)),
         [1, 1, 5, 5],
         [1, 5, 2, 2],
     ),
     (
         "input_grid_interpolation_linear_sample_clamp",
-        (lambda x, grid: torch.ops.aten.grid_sampler(x, grid, 0, 1, True)),
+        "torch.ops.aten.grid_sampler",
+        (lambda x, grid, op: op(x, grid, 1, 1, True)),
         [1, 1, 5, 5],
         [1, 5, 2, 2],
     ),
     (
         "input_grid_interpolation_linear_sample_reflect",
-        (lambda x, grid: torch.ops.aten.grid_sampler(x, grid, 0, 2, True)),
+        "torch.ops.aten.grid_sampler",
+        (lambda x, grid, op: op(x, grid, 1, 2, True)),
         [1, 1, 5, 5],
         [1, 5, 2, 2],
     ),
     (
         "input_grid_interpolation_cubic_sample_fill",
-        (lambda x, grid: torch.ops.aten.grid_sampler(x, grid, 0, 0, True)),
+        "torch.ops.aten.grid_sampler",
+        (lambda x, grid, op: op(x, grid, 2, 0, True)),
         [1, 1, 5, 5],
         [1, 5, 2, 2],
     ),
     (
         "input_grid_interpolation_cubic_sample_clamp",
-        (lambda x, grid: torch.ops.aten.grid_sampler(x, grid, 0, 1, True)),
+        "torch.ops.aten.grid_sampler",
+        (lambda x, grid, op: op(x, grid, 2, 1, True)),
         [1, 1, 5, 5],
         [1, 5, 2, 2],
     ),
     (
         "input_grid_interpolation_cubic_sample_reflect",
-        (lambda x, grid: torch.ops.aten.grid_sampler(x, grid, 0, 2, True)),
-        [1, 1, 5, 5],
-        [1, 5, 2, 2],
-    ),
-    (
-        "input_grid_interpolation_nearest_sample_fill_2d",
-        (lambda x, grid: torch.ops.aten.grid_sampler_2d(x, grid, 0, 0, True)),
-        [1, 1, 5, 5],
-        [1, 5, 2, 2],
-    ),
-    (
-        "input_grid_interpolation_nearest_sample_clamp_2d",
-        (lambda x, grid: torch.ops.aten.grid_sampler_2d(x, grid, 0, 1, True)),
-        [1, 1, 5, 5],
-        [1, 5, 2, 2],
-    ),
-    (
-        "input_grid_interpolation_nearest_sample_reflect_2d",
-        (lambda x, grid: torch.ops.aten.grid_sampler_2d(x, grid, 0, 2, True)),
-        [1, 1, 5, 5],
-        [1, 5, 2, 2],
-    ),
-    (
-        "input_grid_interpolation_linear_sample_fill_2d",
-        (lambda x, grid: torch.ops.aten.grid_sampler_2d(x, grid, 0, 0, True)),
-        [1, 1, 5, 5],
-        [1, 5, 2, 2],
-    ),
-    (
-        "input_grid_interpolation_linear_sample_clamp_2d",
-        (lambda x, grid: torch.ops.aten.grid_sampler_2d(x, grid, 0, 1, True)),
-        [1, 1, 5, 5],
-        [1, 5, 2, 2],
-    ),
-    (
-        "input_grid_interpolation_linear_sample_reflect_2d",
-        (lambda x, grid: torch.ops.aten.grid_sampler_2d(x, grid, 0, 2, True)),
-        [1, 1, 5, 5],
-        [1, 5, 2, 2],
-    ),
-    (
-        "input_grid_interpolation_cubic_sample_fill_2d",
-        (lambda x, grid: torch.ops.aten.grid_sampler_2d(x, grid, 0, 0, True)),
-        [1, 1, 5, 5],
-        [1, 5, 2, 2],
-    ),
-    (
-        "input_grid_interpolation_cubic_sample_clamp_2d",
-        (lambda x, grid: torch.ops.aten.grid_sampler_2d(x, grid, 0, 1, True)),
-        [1, 1, 5, 5],
-        [1, 5, 2, 2],
-    ),
-    (
-        "input_grid_interpolation_cubic_sample_reflect_2d",
-        (lambda x, grid: torch.ops.aten.grid_sampler_2d(x, grid, 0, 2, True)),
+        "torch.ops.aten.grid_sampler",
+        (lambda x, grid, op: op(x, grid, 2, 2, True)),
         [1, 1, 5, 5],
         [1, 5, 2, 2],
     ),
@@ -126,11 +88,90 @@ class TestGridConverter(DispatchTestCase):
                 grid_sampler_op[1],
                 grid_sampler_op[2],
                 grid_sampler_op[3],
+                grid_sampler_op[4],
+            )
+            for grid_sampler_op in grid_sampler_ops
+        ]
+    )
+    def test_grid(self, _, op_name, op, input_shape, dim_shape):
+        class TestModule(nn.Module):
+            def __init__(self, grid_sampler_op):
+                super().__init__()
+                self.grid_sampler_op = grid_sampler_op
+
+            def forward(self, x):
+                grid = torch.randint(-1, 1, dim_shape, dtype=torch.float32)
+                return self.grid_sampler_op(x, grid, grid_sampler_aten_ops[op_name])
+
+        inputs = [torch.randn(input_shape, dtype=torch.float32)]
+        grid_model = TestModule(op)
+        self.run_test(grid_model, inputs)
+
+    @parameterized.expand(
+        [
+            (
+                grid_sampler_op[0],
+                grid_sampler_op[1] + "_2d",
+                grid_sampler_op[2],
+                grid_sampler_op[3],
+                grid_sampler_op[4],
+            )
+            for grid_sampler_op in grid_sampler_ops
+        ]
+    )
+    def test_grid_2d(self, _, op_name, op, input_shape, dim_shape):
+        class TestModule(nn.Module):
+            def __init__(self, grid_sampler_op):
+                super().__init__()
+                self.grid_sampler_op = grid_sampler_op
+
+            def forward(self, x):
+                grid = torch.randint(-1, 1, dim_shape, dtype=torch.float32)
+                return self.grid_sampler_op(x, grid, grid_sampler_aten_ops[op_name])
+
+        inputs = [torch.randn(input_shape, dtype=torch.float32)]
+        grid_model = TestModule(op)
+        self.run_test(grid_model, inputs)
+
+    @parameterized.expand(
+        [
+            (
+                grid_sampler_op[0],
+                grid_sampler_op[1] + ".default",
+                grid_sampler_op[2],
+                grid_sampler_op[3],
+                grid_sampler_op[4],
+            )
+            for grid_sampler_op in grid_sampler_ops
+        ]
+    )
+    def test_grid_default(self, _, op_name, op, input_shape, dim_shape):
+        class TestModule(nn.Module):
+            def __init__(self, grid_sampler_op):
+                super().__init__()
+                self.grid_sampler_op = grid_sampler_op
+
+            def forward(self, x):
+                grid = torch.randint(-1, 1, dim_shape, dtype=torch.float32)
+                return self.grid_sampler_op(x, grid, grid_sampler_aten_ops[op_name])
+
+        inputs = [torch.randn(input_shape, dtype=torch.float32)]
+        grid_model = TestModule(op)
+        self.run_test(grid_model, inputs)
+
+    @parameterized.expand(
+        [
+            (
+                grid_sampler_op[0],
+                grid_sampler_op[1] + "_2d.default",
+                grid_sampler_op[2],
+                grid_sampler_op[3],
+                grid_sampler_op[4],
             )
             for grid_sampler_op in grid_sampler_ops
         ]
     )
-    def test_grid(self, _, op, input_shape, dim_shape):
+    def test_grid_2d_default(self, _, op_name, op, input_shape, dim_shape):
         class TestModule(nn.Module):
             def __init__(self, grid_sampler_op):
                 super().__init__()
@@ -138,7 +179,7 @@ def __init__(self, grid_sampler_op):
 
             def forward(self, x):
                 grid = torch.randint(-1, 1, dim_shape, dtype=torch.float32)
-                return self.grid_sampler_op(x, grid)
+                return self.grid_sampler_op(x, grid, grid_sampler_aten_ops[op_name])
 
         inputs = [torch.randn(input_shape, dtype=torch.float32)]
         grid_model = TestModule(op)