🐛 [Bug] select_scatter decomposition caused compilation failure

Bug Description

The culprit seems to be https://github.com/pytorch/TensorRT/commit/65833002c58b4f02a37224bb516994f2c97faad3.

holy@HOLYWU:~$ python3 test.py
INFO:torch_tensorrt.dynamo.utils:Using Default Torch-TRT Runtime (as requested by user)
INFO:torch_tensorrt.dynamo.utils:Compilation Settings: CompilationSettings(enabled_precisions={<dtype.f16: 6>}, 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_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, assume_dynamic_shape_support=False, sparse_weights=False, refit=False, engine_capability=<EngineCapability.STANDARD: 1>, num_avg_timing_iters=1, dla_sram_size=1048576, dla_local_dram_size=1073741824, dla_global_dram_size=536870912, dryrun=False, hardware_compatible=False)

DEBUG:torch_tensorrt.dynamo.backend.backends:Pre-AOT Autograd graph:
graph():
    %l_x_ : torch.Tensor [num_users=2] = placeholder[target=L_x_]
    %setitem : [num_users=0] = call_function[target=operator.setitem](args = (%l_x_, (slice(None, None, None), slice(None, None, None), 0, 0), 3.1415926), kwargs = {})
    return (l_x_,)
DEBUG:torch_tensorrt.dynamo.lowering._repair_input_aliasing:Inserted auxiliary clone nodes for placeholders:
graph():
    %l_x_ : torch.Tensor [num_users=1] = placeholder[target=L_x_]
    %clone_default : [num_users=2] = call_function[target=torch.ops.aten.clone.default](args = (%l_x_,), kwargs = {})
    %setitem : [num_users=0] = call_function[target=operator.setitem](args = (%clone_default, (slice(None, None, None), slice(None, None, None), 0, 0), 3.1415926), kwargs = {})
    return (clone_default,)
DEBUG:torch_tensorrt.dynamo.lowering._remove_sym_nodes:Removed SymInt placeholders:
graph():
    %l_x_ : torch.Tensor [num_users=1] = placeholder[target=L_x_]
    %clone_default : [num_users=2] = call_function[target=torch.ops.aten.clone.default](args = (%l_x_,), kwargs = {})
    %setitem : [num_users=0] = call_function[target=operator.setitem](args = (%clone_default, (slice(None, None, None), slice(None, None, None), 0, 0), 3.1415926), kwargs = {})
    return (clone_default,)
DEBUG:torch_tensorrt.dynamo.backend.backends:Post-AOT Autograd graph:
graph():
    %arg0_1 : [num_users=1] = placeholder[target=arg0_1]
    %clone : [num_users=2] = call_function[target=torch.ops.aten.clone.default](args = (%arg0_1,), kwargs = {})
    %_tensor_constant0 : [num_users=1] = get_attr[target=_tensor_constant0]
    %lift_fresh_copy : [num_users=1] = call_function[target=torch.ops.aten.lift_fresh_copy.default](args = (%_tensor_constant0,), kwargs = {})
    %slice_1 : [num_users=1] = call_function[target=torch.ops.aten.slice.Tensor](args = (%clone, 0, 0, 9223372036854775807), kwargs = {})
    %slice_2 : [num_users=1] = call_function[target=torch.ops.aten.slice.Tensor](args = (%slice_1, 1, 0, 9223372036854775807), kwargs = {})
    %select : [num_users=1] = call_function[target=torch.ops.aten.select.int](args = (%slice_2, 2, 0), kwargs = {})
    %select_1 : [num_users=1] = call_function[target=torch.ops.aten.select.int](args = (%select, 2, 0), kwargs = {})
    %copy : [num_users=1] = call_function[target=torch.ops.aten.copy.default](args = (%select_1, %lift_fresh_copy), kwargs = {})
    %slice_3 : [num_users=1] = call_function[target=torch.ops.aten.slice.Tensor](args = (%clone, 0, 0, 9223372036854775807), kwargs = {})
    %slice_4 : [num_users=2] = call_function[target=torch.ops.aten.slice.Tensor](args = (%slice_3, 1, 0, 9223372036854775807), kwargs = {})
    %select_2 : [num_users=1] = call_function[target=torch.ops.aten.select.int](args = (%slice_4, 2, 0), kwargs = {})
    %unsqueeze : [num_users=1] = call_function[target=torch.ops.aten.unsqueeze.default](args = (%copy, 2), kwargs = {})
    %empty_strided : [num_users=1] = call_function[target=torch.ops.aten.empty_strided.default](args = ([1, 1], [1, 1]), kwargs = {dtype: torch.int64, layout: torch.strided, device: cpu, pin_memory: False})
    %full_like : [num_users=1] = call_function[target=torch.ops.aten.full_like.default](args = (%empty_strided, 1), kwargs = {pin_memory: False})
    %mul : [num_users=1] = call_function[target=torch.ops.aten.mul.Tensor](args = (%full_like, 0), kwargs = {})
    %unsqueeze_1 : [num_users=1] = call_function[target=torch.ops.aten.unsqueeze.default](args = (%mul, 2), kwargs = {})
    %cat : [num_users=1] = call_function[target=torch.ops.aten.cat.default](args = ([%unsqueeze_1], 2), kwargs = {})
    %_to_copy : [num_users=1] = call_function[target=torch.ops.aten._to_copy.default](args = (%cat,), kwargs = {dtype: torch.int64, layout: torch.strided, device: cuda:0})
    %scatter : [num_users=1] = call_function[target=torch.ops.aten.scatter.src](args = (%select_2, 2, %_to_copy, %unsqueeze), kwargs = {})
    %unsqueeze_2 : [num_users=1] = call_function[target=torch.ops.aten.unsqueeze.default](args = (%scatter, 2), kwargs = {})
    %empty_strided_1 : [num_users=1] = call_function[target=torch.ops.aten.empty_strided.default](args = ([1, 1, 4], [4, 4, 1]), kwargs = {dtype: torch.int64, layout: torch.strided, device: cpu, pin_memory: False})
    %full_like_1 : [num_users=1] = call_function[target=torch.ops.aten.full_like.default](args = (%empty_strided_1, 1), kwargs = {pin_memory: False})
    %mul_1 : [num_users=1] = call_function[target=torch.ops.aten.mul.Tensor](args = (%full_like_1, 0), kwargs = {})
    %cat_1 : [num_users=1] = call_function[target=torch.ops.aten.cat.default](args = ([%mul_1], 2), kwargs = {})
    %view : [num_users=1] = call_function[target=torch.ops.aten.view.default](args = (%cat_1, [1, 1, 1, 4]), kwargs = {})
    %_to_copy_1 : [num_users=1] = call_function[target=torch.ops.aten._to_copy.default](args = (%view,), kwargs = {dtype: torch.int64, layout: torch.strided, device: cuda:0})
    %scatter_1 : [num_users=1] = call_function[target=torch.ops.aten.scatter.src](args = (%slice_4, 2, %_to_copy_1, %unsqueeze_2), kwargs = {})
    return (scatter_1,)
DEBUG:torch_tensorrt.dynamo.lowering.passes.remove_input_alias_fixing_clones:Removing node clone from graph, since it is a clone node which is the only user of placeholder arg0_1 and was inserted by the compiler.
DEBUG:torch_tensorrt.dynamo.lowering.passes.remove_input_alias_fixing_clones:Removed auxiliary clone nodes for placeholders:
graph():
    %arg0_1 : [num_users=2] = placeholder[target=arg0_1]
    %_tensor_constant0 : [num_users=1] = get_attr[target=_tensor_constant0]
    %lift_fresh_copy : [num_users=1] = call_function[target=torch.ops.aten.lift_fresh_copy.default](args = (%_tensor_constant0,), kwargs = {})
    %slice_1 : [num_users=1] = call_function[target=torch.ops.aten.slice.Tensor](args = (%arg0_1, 0, 0, 9223372036854775807), kwargs = {})
    %slice_2 : [num_users=1] = call_function[target=torch.ops.aten.slice.Tensor](args = (%slice_1, 1, 0, 9223372036854775807), kwargs = {})
    %select : [num_users=1] = call_function[target=torch.ops.aten.select.int](args = (%slice_2, 2, 0), kwargs = {})
    %select_1 : [num_users=1] = call_function[target=torch.ops.aten.select.int](args = (%select, 2, 0), kwargs = {})
    %copy : [num_users=1] = call_function[target=torch.ops.aten.copy.default](args = (%select_1, %lift_fresh_copy), kwargs = {})
    %slice_3 : [num_users=1] = call_function[target=torch.ops.aten.slice.Tensor](args = (%arg0_1, 0, 0, 9223372036854775807), kwargs = {})
    %slice_4 : [num_users=2] = call_function[target=torch.ops.aten.slice.Tensor](args = (%slice_3, 1, 0, 9223372036854775807), kwargs = {})
    %select_2 : [num_users=1] = call_function[target=torch.ops.aten.select.int](args = (%slice_4, 2, 0), kwargs = {})
    %unsqueeze : [num_users=1] = call_function[target=torch.ops.aten.unsqueeze.default](args = (%copy, 2), kwargs = {})
    %empty_strided : [num_users=1] = call_function[target=torch.ops.aten.empty_strided.default](args = ([1, 1], [1, 1]), kwargs = {dtype: torch.int64, layout: torch.strided, device: cpu, pin_memory: False})
    %full_like : [num_users=1] = call_function[target=torch.ops.aten.full_like.default](args = (%empty_strided, 1), kwargs = {pin_memory: False})
    %mul : [num_users=1] = call_function[target=torch.ops.aten.mul.Tensor](args = (%full_like, 0), kwargs = {})
    %unsqueeze_1 : [num_users=1] = call_function[target=torch.ops.aten.unsqueeze.default](args = (%mul, 2), kwargs = {})
    %cat : [num_users=1] = call_function[target=torch.ops.aten.cat.default](args = ([%unsqueeze_1], 2), kwargs = {})
    %_to_copy : [num_users=1] = call_function[target=torch.ops.aten._to_copy.default](args = (%cat,), kwargs = {dtype: torch.int64, layout: torch.strided, device: cuda:0})
    %scatter : [num_users=1] = call_function[target=torch.ops.aten.scatter.src](args = (%select_2, 2, %_to_copy, %unsqueeze), kwargs = {})
    %unsqueeze_2 : [num_users=1] = call_function[target=torch.ops.aten.unsqueeze.default](args = (%scatter, 2), kwargs = {})
    %empty_strided_1 : [num_users=1] = call_function[target=torch.ops.aten.empty_strided.default](args = ([1, 1, 4], [4, 4, 1]), kwargs = {dtype: torch.int64, layout: torch.strided, device: cpu, pin_memory: False})
    %full_like_1 : [num_users=1] = call_function[target=torch.ops.aten.full_like.default](args = (%empty_strided_1, 1), kwargs = {pin_memory: False})
    %mul_1 : [num_users=1] = call_function[target=torch.ops.aten.mul.Tensor](args = (%full_like_1, 0), kwargs = {})
    %cat_1 : [num_users=1] = call_function[target=torch.ops.aten.cat.default](args = ([%mul_1], 2), kwargs = {})
    %view : [num_users=1] = call_function[target=torch.ops.aten.view.default](args = (%cat_1, [1, 1, 1, 4]), kwargs = {})
    %_to_copy_1 : [num_users=1] = call_function[target=torch.ops.aten._to_copy.default](args = (%view,), kwargs = {dtype: torch.int64, layout: torch.strided, device: cuda:0})
    %scatter_1 : [num_users=1] = call_function[target=torch.ops.aten.scatter.src](args = (%slice_4, 2, %_to_copy_1, %unsqueeze_2), kwargs = {})
    return (scatter_1,)
DEBUG:torch_tensorrt.dynamo.lowering.passes.constant_folding:Graph after constant folding:
graph():
    %arg0_1 : [num_users=2] = placeholder[target=arg0_1]
    %_frozen_param0 : [num_users=1] = get_attr[target=_frozen_param0]
    %slice_1 : [num_users=1] = call_function[target=torch.ops.aten.slice.Tensor](args = (%arg0_1, 0, 0, 9223372036854775807), kwargs = {})
    %slice_2 : [num_users=1] = call_function[target=torch.ops.aten.slice.Tensor](args = (%slice_1, 1, 0, 9223372036854775807), kwargs = {})
    %select : [num_users=1] = call_function[target=torch.ops.aten.select.int](args = (%slice_2, 2, 0), kwargs = {})
    %select_1 : [num_users=1] = call_function[target=torch.ops.aten.select.int](args = (%select, 2, 0), kwargs = {})
    %copy : [num_users=1] = call_function[target=torch.ops.aten.copy.default](args = (%select_1, %_frozen_param0), kwargs = {})
    %slice_3 : [num_users=1] = call_function[target=torch.ops.aten.slice.Tensor](args = (%arg0_1, 0, 0, 9223372036854775807), kwargs = {})
    %slice_4 : [num_users=2] = call_function[target=torch.ops.aten.slice.Tensor](args = (%slice_3, 1, 0, 9223372036854775807), kwargs = {})
    %select_2 : [num_users=1] = call_function[target=torch.ops.aten.select.int](args = (%slice_4, 2, 0), kwargs = {})
    %unsqueeze : [num_users=1] = call_function[target=torch.ops.aten.unsqueeze.default](args = (%copy, 2), kwargs = {})
    %_frozen_param1 : [num_users=1] = get_attr[target=_frozen_param1]
    %scatter : [num_users=1] = call_function[target=torch.ops.aten.scatter.src](args = (%select_2, 2, %_frozen_param1, %unsqueeze), kwargs = {})
    %unsqueeze_2 : [num_users=1] = call_function[target=torch.ops.aten.unsqueeze.default](args = (%scatter, 2), kwargs = {})
    %_frozen_param2 : [num_users=1] = get_attr[target=_frozen_param2]
    %scatter_1 : [num_users=1] = call_function[target=torch.ops.aten.scatter.src](args = (%slice_4, 2, %_frozen_param2, %unsqueeze_2), kwargs = {})
    return (scatter_1,)
DEBUG:torch_tensorrt.dynamo.partitioning._global_partitioner:
Supported Nodes:
- torch.ops.aten.slice.Tensor + Operator Count: 4
- torch.ops.aten.select.int + Operator Count: 3
- torch.ops.aten.copy.default + Operator Count: 1
- torch.ops.aten.unsqueeze.default + Operator Count: 2
- torch.ops.aten.scatter.src + Operator Count: 2

DEBUG:torch_tensorrt.dynamo.partitioning._global_partitioner:
All Nodes Supported

DEBUG:torch_tensorrt.dynamo._compiler:Detected support for 12 operators out of 12 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.slice.Tensor + Operator Count: 4
- torch.ops.aten.select.int + Operator Count: 3
- torch.ops.aten.copy.default + Operator Count: 1
- torch.ops.aten.unsqueeze.default + Operator Count: 2
- torch.ops.aten.scatter.src + Operator Count: 2

DEBUG:torch_tensorrt.dynamo.partitioning._adjacency_partitioner:
All Nodes Supported

DEBUG:torch_tensorrt.dynamo._compiler:Submodule name: _run_on_acc_0
 Input shapes: [(1, 1, 4, 4)]
 graph():
    %arg0_1 : [num_users=2] = placeholder[target=arg0_1]
    %slice_1 : [num_users=1] = call_function[target=torch.ops.aten.slice.Tensor](args = (%arg0_1, 0, 0, 9223372036854775807), kwargs = {})
    %slice_2 : [num_users=1] = call_function[target=torch.ops.aten.slice.Tensor](args = (%slice_1, 1, 0, 9223372036854775807), kwargs = {})
    %select : [num_users=1] = call_function[target=torch.ops.aten.select.int](args = (%slice_2, 2, 0), kwargs = {})
    %select_1 : [num_users=1] = call_function[target=torch.ops.aten.select.int](args = (%select, 2, 0), kwargs = {})
    %_frozen_param0 : [num_users=1] = get_attr[target=_frozen_param0]
    %copy : [num_users=1] = call_function[target=torch.ops.aten.copy.default](args = (%select_1, %_frozen_param0), kwargs = {})
    %slice_3 : [num_users=1] = call_function[target=torch.ops.aten.slice.Tensor](args = (%arg0_1, 0, 0, 9223372036854775807), kwargs = {})
    %slice_4 : [num_users=2] = call_function[target=torch.ops.aten.slice.Tensor](args = (%slice_3, 1, 0, 9223372036854775807), kwargs = {})
    %select_2 : [num_users=1] = call_function[target=torch.ops.aten.select.int](args = (%slice_4, 2, 0), kwargs = {})
    %unsqueeze : [num_users=1] = call_function[target=torch.ops.aten.unsqueeze.default](args = (%copy, 2), kwargs = {})
    %_frozen_param1 : [num_users=1] = get_attr[target=_frozen_param1]
    %scatter : [num_users=1] = call_function[target=torch.ops.aten.scatter.src](args = (%select_2, 2, %_frozen_param1, %unsqueeze), kwargs = {})
    %unsqueeze_2 : [num_users=1] = call_function[target=torch.ops.aten.unsqueeze.default](args = (%scatter, 2), kwargs = {})
    %_frozen_param2 : [num_users=1] = get_attr[target=_frozen_param2]
    %scatter_1 : [num_users=1] = call_function[target=torch.ops.aten.scatter.src](args = (%slice_4, 2, %_frozen_param2, %unsqueeze_2), kwargs = {})
    return scatter_1
INFO:torch_tensorrt [TensorRT Conversion Context]:[MemUsageChange] Init CUDA: CPU +2, GPU +0, now: CPU 141, GPU 1009 (MiB)
INFO:torch_tensorrt [TensorRT Conversion Context]:[MemUsageChange] Init builder kernel library: CPU +1764, GPU +310, now: CPU 2041, GPU 1319 (MiB)
DEBUG:torch_tensorrt.dynamo.conversion._TRTInterpreter:Graph to be compiled to TensorRT: graph():
    %arg0_1 : [num_users=2] = placeholder[target=arg0_1]
    %slice_1 : [num_users=1] = call_function[target=torch.ops.aten.slice.Tensor](args = (%arg0_1, 0, 0, 9223372036854775807), kwargs = {})
    %slice_2 : [num_users=1] = call_function[target=torch.ops.aten.slice.Tensor](args = (%slice_1, 1, 0, 9223372036854775807), kwargs = {})
    %select : [num_users=1] = call_function[target=torch.ops.aten.select.int](args = (%slice_2, 2, 0), kwargs = {})
    %select_1 : [num_users=1] = call_function[target=torch.ops.aten.select.int](args = (%select, 2, 0), kwargs = {})
    %_frozen_param0 : [num_users=1] = get_attr[target=_frozen_param0]
    %copy : [num_users=1] = call_function[target=torch.ops.aten.copy.default](args = (%select_1, %_frozen_param0), kwargs = {})
    %slice_3 : [num_users=1] = call_function[target=torch.ops.aten.slice.Tensor](args = (%arg0_1, 0, 0, 9223372036854775807), kwargs = {})
    %slice_4 : [num_users=2] = call_function[target=torch.ops.aten.slice.Tensor](args = (%slice_3, 1, 0, 9223372036854775807), kwargs = {})
    %select_2 : [num_users=1] = call_function[target=torch.ops.aten.select.int](args = (%slice_4, 2, 0), kwargs = {})
    %unsqueeze : [num_users=1] = call_function[target=torch.ops.aten.unsqueeze.default](args = (%copy, 2), kwargs = {})
    %_frozen_param1 : [num_users=1] = get_attr[target=_frozen_param1]
    %scatter : [num_users=1] = call_function[target=torch.ops.aten.scatter.src](args = (%select_2, 2, %_frozen_param1, %unsqueeze), kwargs = {})
    %unsqueeze_2 : [num_users=1] = call_function[target=torch.ops.aten.unsqueeze.default](args = (%scatter, 2), kwargs = {})
    %_frozen_param2 : [num_users=1] = get_attr[target=_frozen_param2]
    %scatter_1 : [num_users=1] = call_function[target=torch.ops.aten.scatter.src](args = (%slice_4, 2, %_frozen_param2, %unsqueeze_2), kwargs = {})
    return scatter_1
DEBUG:torch_tensorrt.dynamo.conversion._TRTInterpreter:Adding input to in-progress INetwork: arg0_1 [shape=[1, 1, 4, 4], dtype=DataType.HALF]
DEBUG:torch_tensorrt.dynamo.conversion._TRTInterpreter:Converting node slice_1 (kind: aten.slice.Tensor, args: ('arg0_1 <tensorrt.ITensor [shape=(1, 1, 4, 4), dtype=DataType.HALF]>', 0, 0, 9223372036854775807))
DEBUG:torch_tensorrt.dynamo.conversion._TRTInterpreter:Converting node slice_2 (kind: aten.slice.Tensor, args: ('[SLICE]-[unknown_ir_ops.slice.Tensor]-[slice_1]_output <tensorrt.ITensor [shape=(1, 1, 4, 4), dtype=DataType.HALF]>', 1, 0, 9223372036854775807))
DEBUG:torch_tensorrt.dynamo.conversion._TRTInterpreter:Converting node select (kind: aten.select.int, args: ('[SLICE]-[unknown_ir_ops.slice.Tensor]-[slice_2]_output <tensorrt.ITensor [shape=(1, 1, 4, 4), dtype=DataType.HALF]>', 2, 0))
DEBUG:torch_tensorrt.dynamo.conversion._TRTInterpreter:Converting node select_1 (kind: aten.select.int, args: ('(Unnamed Layer* 4) [Shuffle]_output <tensorrt.ITensor [shape=(1, 1, 4), dtype=DataType.HALF]>', 2, 0))
DEBUG:torch_tensorrt.dynamo.conversion._TRTInterpreter:Converting node copy (kind: aten.copy.default, args: ('(Unnamed Layer* 7) [Shuffle]_output <tensorrt.ITensor [shape=(1, 1), dtype=DataType.HALF]>', '<torch.Tensor as np.ndarray [shape=(), dtype=float16]>'))
DEBUG:torch_tensorrt.dynamo.conversion.converter_utils:Freezing tensor copy_constant_1 to TRT IConstantLayer
DEBUG:torch_tensorrt.dynamo.conversion._TRTInterpreter:Converting node slice_3 (kind: aten.slice.Tensor, args: ('arg0_1 <tensorrt.ITensor [shape=(1, 1, 4, 4), dtype=DataType.HALF]>', 0, 0, 9223372036854775807))
DEBUG:torch_tensorrt.dynamo.conversion._TRTInterpreter:Converting node slice_4 (kind: aten.slice.Tensor, args: ('[SLICE]-[unknown_ir_ops.slice.Tensor]-[slice_3]_output <tensorrt.ITensor [shape=(1, 1, 4, 4), dtype=DataType.HALF]>', 1, 0, 9223372036854775807))
DEBUG:torch_tensorrt.dynamo.conversion._TRTInterpreter:Converting node select_2 (kind: aten.select.int, args: ('[SLICE]-[unknown_ir_ops.slice.Tensor]-[slice_4]_output <tensorrt.ITensor [shape=(1, 1, 4, 4), dtype=DataType.HALF]>', 2, 0))
DEBUG:torch_tensorrt.dynamo.conversion._TRTInterpreter:Converting node unsqueeze (kind: aten.unsqueeze.default, args: ('Forced Cast ITensor copy_constant_1_output from DataType.HALF to DataType.HALF - [aten_ops.torch.ops.aten.copy.default]-[copy]_output <tensorrt.ITensor [shape=(), dtype=DataType.HALF]>', 2))
DEBUG:torch_tensorrt.dynamo.conversion._TRTInterpreter:Converting node scatter (kind: aten.scatter.src, args: ('(Unnamed Layer* 14) [Shuffle]_output <tensorrt.ITensor [shape=(1, 1, 4), dtype=DataType.HALF]>', 2, '<torch.Tensor as np.ndarray [shape=(1, 1, 1), dtype=int64]>', '[SHUFFLE]-[aten_ops.unsqueeze.default]-[unsqueeze]_output <tensorrt.ITensor [shape=(1,), dtype=DataType.HALF]>'))
DEBUG:torch_tensorrt.dynamo.conversion.converter_utils:Freezing tensor scatter_constant_2 to TRT IConstantLayer
ERROR:torch_tensorrt [TensorRT Conversion Context]:4: [SCATTER]-[aten_ops.scatter.src]-[scatter_scatter_layer]: ScatterLayer in elements mode all inputs tensors rank must be same.
DEBUG:torch_tensorrt.dynamo.conversion._TRTInterpreter:Converting node unsqueeze_2 (kind: aten.unsqueeze.default, args: ('[SCATTER]-[aten_ops.scatter.src]-[scatter_scatter_layer]_output <tensorrt.ITensor [shape=(1, 1, 4), dtype=DataType.HALF]>', 2))
ERROR:torch_tensorrt [TensorRT Conversion Context]:4: [graphShapeAnalyzer.cpp::needTypeAndDimensions::2276] Error Code 4: Internal Error ([SCATTER]-[aten_ops.scatter.src]-[scatter_scatter_layer]: output shape can not be computed)
Traceback (most recent call last):
  File "/home/holy/test.py", line 30, in <module>
    print(optimized_model(*inputs))
  File "/home/holy/.local/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1545, in _wrapped_call_impl
    return self._call_impl(*args, **kwargs)
  File "/home/holy/.local/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1554, in _call_impl
    return forward_call(*args, **kwargs)
  File "/home/holy/.local/lib/python3.10/site-packages/torch/_dynamo/eval_frame.py", line 421, in _fn
    return fn(*args, **kwargs)
  File "/home/holy/.local/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1545, in _wrapped_call_impl
    return self._call_impl(*args, **kwargs)
  File "/home/holy/.local/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1554, in _call_impl
    return forward_call(*args, **kwargs)
  File "/home/holy/.local/lib/python3.10/site-packages/torch/_dynamo/convert_frame.py", line 1077, in catch_errors
    return callback(frame, cache_entry, hooks, frame_state, skip=1)
  File "/home/holy/.local/lib/python3.10/site-packages/torch/_dynamo/convert_frame.py", line 918, in _convert_frame
    result = inner_convert(
  File "/home/holy/.local/lib/python3.10/site-packages/torch/_dynamo/convert_frame.py", line 456, in _convert_frame_assert
    return _compile(
  File "/home/holy/.local/lib/python3.10/site-packages/torch/_utils_internal.py", line 83, in wrapper_function
    return StrobelightCompileTimeProfiler.profile_compile_time(
  File "/home/holy/.local/lib/python3.10/site-packages/torch/_strobelight/compile_time_profiler.py", line 129, in profile_compile_time
    return func(*args, **kwargs)
  File "/usr/lib/python3.10/contextlib.py", line 79, in inner
    return func(*args, **kwds)
  File "/home/holy/.local/lib/python3.10/site-packages/torch/_dynamo/convert_frame.py", line 799, in _compile
    guarded_code = compile_inner(code, one_graph, hooks, transform)
  File "/home/holy/.local/lib/python3.10/site-packages/torch/_dynamo/utils.py", line 218, in time_wrapper
    r = func(*args, **kwargs)
  File "/home/holy/.local/lib/python3.10/site-packages/torch/_dynamo/convert_frame.py", line 618, in compile_inner
    out_code = transform_code_object(code, transform)
  File "/home/holy/.local/lib/python3.10/site-packages/torch/_dynamo/bytecode_transformation.py", line 1184, in transform_code_object
    transformations(instructions, code_options)
  File "/home/holy/.local/lib/python3.10/site-packages/torch/_dynamo/convert_frame.py", line 177, in _fn
    return fn(*args, **kwargs)
  File "/home/holy/.local/lib/python3.10/site-packages/torch/_dynamo/convert_frame.py", line 564, in transform
    tracer.run()
  File "/home/holy/.local/lib/python3.10/site-packages/torch/_dynamo/symbolic_convert.py", line 2248, in run
    super().run()
  File "/home/holy/.local/lib/python3.10/site-packages/torch/_dynamo/symbolic_convert.py", line 888, in run
    while self.step():
  File "/home/holy/.local/lib/python3.10/site-packages/torch/_dynamo/symbolic_convert.py", line 803, in step
    self.dispatch_table[inst.opcode](self, inst)
  File "/home/holy/.local/lib/python3.10/site-packages/torch/_dynamo/symbolic_convert.py", line 2439, in RETURN_VALUE
    self._return(inst)
  File "/home/holy/.local/lib/python3.10/site-packages/torch/_dynamo/symbolic_convert.py", line 2424, in _return
    self.output.compile_subgraph(
  File "/home/holy/.local/lib/python3.10/site-packages/torch/_dynamo/output_graph.py", line 1097, in compile_subgraph
    self.compile_and_call_fx_graph(tx, list(reversed(stack_values)), root)
  File "/usr/lib/python3.10/contextlib.py", line 79, in inner
    return func(*args, **kwds)
  File "/home/holy/.local/lib/python3.10/site-packages/torch/_dynamo/output_graph.py", line 1314, in compile_and_call_fx_graph
    compiled_fn = self.call_user_compiler(gm)
  File "/home/holy/.local/lib/python3.10/site-packages/torch/_dynamo/utils.py", line 218, in time_wrapper
    r = func(*args, **kwargs)
  File "/home/holy/.local/lib/python3.10/site-packages/torch/_dynamo/output_graph.py", line 1405, in call_user_compiler
    raise BackendCompilerFailed(self.compiler_fn, e).with_traceback(
  File "/home/holy/.local/lib/python3.10/site-packages/torch/_dynamo/output_graph.py", line 1386, in call_user_compiler
    compiled_fn = compiler_fn(gm, self.example_inputs())
  File "/home/holy/.local/lib/python3.10/site-packages/torch/_dynamo/repro/after_dynamo.py", line 127, in debug_wrapper
    compiled_gm = compiler_fn(gm, example_inputs)
  File "/home/holy/.local/lib/python3.10/site-packages/torch/__init__.py", line 1831, in __call__
    return self.compiler_fn(model_, inputs_, **self.kwargs)
  File "/home/holy/.local/lib/python3.10/site-packages/torch_tensorrt/dynamo/backend/backends.py", line 43, in torch_tensorrt_backend
    return DEFAULT_BACKEND(gm, sample_inputs, **kwargs)
  File "/home/holy/.local/lib/python3.10/site-packages/torch_tensorrt/dynamo/backend/backends.py", line 51, in aot_torch_tensorrt_aten_backend
    return _pretraced_backend(gm, sample_inputs, settings)
  File "/home/holy/.local/lib/python3.10/site-packages/torch_tensorrt/dynamo/backend/backends.py", line 102, in _pretraced_backend
    trt_compiled = compile_module(
  File "/home/holy/.local/lib/python3.10/site-packages/torch_tensorrt/dynamo/_compiler.py", line 408, in compile_module
    trt_module = convert_module(
  File "/home/holy/.local/lib/python3.10/site-packages/torch_tensorrt/dynamo/conversion/_conversion.py", line 106, in convert_module
    interpreter_result = interpret_module_to_result(module, inputs, settings)
  File "/home/holy/.local/lib/python3.10/site-packages/torch_tensorrt/dynamo/conversion/_conversion.py", line 87, in interpret_module_to_result
    interpreter_result = interpreter.run()
  File "/home/holy/.local/lib/python3.10/site-packages/torch_tensorrt/dynamo/conversion/_TRTInterpreter.py", line 308, in run
    super().run()
  File "/home/holy/.local/lib/python3.10/site-packages/torch/fx/interpreter.py", line 145, in run
    self.env[node] = self.run_node(node)
  File "/home/holy/.local/lib/python3.10/site-packages/torch_tensorrt/dynamo/conversion/_TRTInterpreter.py", line 347, in run_node
    trt_node: torch.fx.Node = super().run_node(n)
  File "/home/holy/.local/lib/python3.10/site-packages/torch/fx/interpreter.py", line 202, in run_node
    return getattr(self, n.op)(n.target, args, kwargs)
  File "/home/holy/.local/lib/python3.10/site-packages/torch_tensorrt/dynamo/conversion/_TRTInterpreter.py", line 443, in call_function
    return converter(self.ctx, target, args, kwargs, self._cur_node_name)
  File "/home/holy/.local/lib/python3.10/site-packages/torch_tensorrt/dynamo/conversion/aten_ops_converters.py", line 613, in aten_ops_unsqueeze
    return impl.unsqueeze.unsqueeze(
  File "/home/holy/.local/lib/python3.10/site-packages/torch_tensorrt/dynamo/conversion/impl/unsqueeze.py", line 32, in unsqueeze
    input_shape_size = len(input_val.shape)
torch._dynamo.exc.BackendCompilerFailed: backend='torch_tensorrt_backend' raised:
ValueError: __len__() should return >= 0

While executing %unsqueeze_2 : [num_users=1] = call_function[target=torch.ops.aten.unsqueeze.default](args = (%scatter, 2), kwargs = {_itensor_to_tensor_meta: {<tensorrt_bindings.tensorrt.ITensor object at 0x7fd623671cf0>: ((1, 1, 4, 4), torch.float16, False, (16, 16, 4, 1), torch.contiguous_format, False, {}), <tensorrt_bindings.tensorrt.ITensor object at 0x7fd6236681f0>: ((1, 1, 4, 4), torch.float16, False, (16, 16, 4, 1), torch.contiguous_format, False, {}), <tensorrt_bindings.tensorrt.ITensor object at 0x7fd6236720f0>: ((1, 1, 4, 4), torch.float16, False, (16, 16, 4, 1), torch.contiguous_format, False, {}), <tensorrt_bindings.tensorrt.ITensor object at 0x7fd6236722b0>: ((1, 1, 4), torch.float16, False, (16, 16, 1), torch.contiguous_format, False, {}), <tensorrt_bindings.tensorrt.ITensor object at 0x7fd6236726f0>: ((1, 1), torch.float16, False, (16, 16), torch.contiguous_format, False, {}), <tensorrt_bindings.tensorrt.ITensor object at 0x7fd623673130>: ((1, 1), torch.float16, False, (16, 16), torch.contiguous_format, False, {}), <tensorrt_bindings.tensorrt.ITensor object at 0x7fd623673670>: ((1, 1, 4, 4), torch.float16, False, (16, 16, 4, 1), torch.contiguous_format, False, {}), <tensorrt_bindings.tensorrt.ITensor object at 0x7fd623673cb0>: ((1, 1, 4, 4), torch.float16, False, (16, 16, 4, 1), torch.contiguous_format, False, {}), <tensorrt_bindings.tensorrt.ITensor object at 0x7fd6236881b0>: ((1, 1, 4), torch.float16, False, (16, 16, 1), torch.contiguous_format, False, {}), <tensorrt_bindings.tensorrt.ITensor object at 0x7fd623688b70>: ((1, 1, 1), torch.float16, False, (16, 16, 1), torch.contiguous_format, False, {}), <tensorrt_bindings.tensorrt.ITensor object at 0x7fd623689cf0>: ((1, 1, 4), torch.float16, False, (4, 4, 1), torch.contiguous_format, False, {})}})
Original traceback:
None

Set TORCH_LOGS="+dynamo" and TORCHDYNAMO_VERBOSE=1 for more information


You can suppress this exception and fall back to eager by setting:
    import torch._dynamo
    torch._dynamo.config.suppress_errors = True

To Reproduce

import torch
import torch.nn as nn
import torch_tensorrt


class MyModule(nn.Module):
    def __init__(self):
        super().__init__()

    def forward(self, x):
        x[:, :, 0, 0] = 3.1415926
        return x


device = torch.device("cuda", 0)
model = MyModule().eval().to(device).half()
inputs = [torch.zeros((1, 1, 4, 4), dtype=torch.half, device=device)]

optimized_model = torch_tensorrt.compile(
    model,
    ir="torch_compile",
    inputs=inputs,
    enabled_precisions={torch.half},
    debug=True,
    min_block_size=1,
    truncate_double=True,
    device=device,
)

print(optimized_model(*inputs))

Environment

• Torch-TensorRT Version (e.g. 1.0.0): 2.4.0.dev20240601+cu124
• PyTorch Version (e.g. 1.0): 2.4.0.dev20240601+cu124
• CPU Architecture: x64
• OS (e.g., Linux): Ubuntu 22.04
• 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.4
• GPU models and configuration: RTX 3050
• Any other relevant information:

Additional context

Although https://github.com/pytorch/TensorRT/blob/91637208e2e81ceac890a7f334ab535f695cabca/py/torch_tensorrt/dynamo/lowering/_decomposition_groups.py#L120 has nothing to do with this issue, this line should still be removed in favor of Torch-TRT's own decomposition.