Repository: pytorch/TensorRT

Issue:
🐛 [Bug] Encountered bug when compiling model with nested Input type using `input_signature` argument
##  Bug Description

When compiling a model with input type `Tuple[torch.Tensor, Tuple[torch.Tensor, torch.Tensor]]`, compilation fails with the following error:
```python
RuntimeError: [Error thrown at core/ir/ir.cpp:46] Expected vals.size() == specs.size() to be true but got false
Expected dimension specifications for all input tensors, but found 1 input tensors and 2 dimension specs
```

## To Reproduce

Steps to reproduce the behavior:

1. Define a Torch model with forward function having the following form:
```python
def forward(self, input : Tuple[torch.Tensor, Tuple[torch.Tensor, torch.Tensor]])
```
2. Define the Torch-TRT Inputs and compilation settings using `input_signature`, then compile the scripted model:
```python
compile_settings = {
    "input_signature": (torch_tensorrt.Input((5, 5), dtype=torch.float), (torch_tensorrt.Input((5, 5), dtype=torch.float), torch_tensorrt.Input((5, 5), dtype=torch.float)),),
    "enabled_precisions": {torch.float},
    "truncate_long_and_double": True,
}
trt_ts_module = torch_tensorrt.ts.compile(scripted_model, **compile_settings)
```

## Expected behavior

Model should compile with input signature containing nested Tuple collection.

## Environment

- Torch-TensorRT Version: 1.4.0.dev0+f43be5b6
- PyTorch Version: 1.14.0.dev20221114+cu116
- CPU Architecture: Intel Xeon CPU
- OS: Ubuntu 20.04
- How you installed PyTorch: pip
- Build command you used: `python setup.py develop`
- Are you using local sources or building from archives: local
- Python version: 3.8.13
- CUDA version: 11.6

## Additional context

The case of a nested Tuple containing a singleton Tensor followed by a Tuple of Tensors is not parsed as expected by the input signature interpreter, whereas regular tuple Inputs of Tensors are supported, as in:
https://github.com/pytorch/TensorRT/blob/b2a5da66cd22db22560bdb7b54f107fb4bcffb2e/tests/py/api/test_collections.py#L210-L213

Wrapping the input signature in another Tuple does not resolve the issue, and raises a different error message:
```python
RuntimeError: forward() Expected a value of type 'Tuple[Tensor, Tuple[Tensor, Tensor]]' for argument 'input.1' but instead found type 'Tuple[Tensor]'.
Position: 1
Declaration: forward.forward(__torch__.module self_1, (Tensor, (Tensor, Tensor)) input.1) -> ((Tensor, Tensor, Tensor) 7)
```


Discussion:
## Update
### Error Origin
The origin of this compilation error derives from the graph inputs sourced from the `input_signature` argument. As per the documentation of `input_signature`, the signature is expected as a Python Tuple or List of input arguments to the model's `forward` function. This implies that the desired format of the inputs for the above case would be:
```python
( (Tensor, (Tensor, Tensor)), )
```
With the above as the input format, however, the innermost `Tensor` object is considered to have nesting depth 3, which is disallowed here:
https://github.com/pytorch/TensorRT/blob/5fa63743fc4448e38b1e6b01e7403c0d74717158/core/ir/GraphInputs.cpp#L42-L52
Since the initial Tuple coalescing the individual arguments in the input signature is considered as a nesting level, this implies that `Tuple[Tuple[Tensor]]` input constructs are disallowed as they violate the maximum nesting depth.

### Resolution
Since this example did not derive from a specific Input, but was instead a general assessment of `input_signature`, resolving the error for this specific signature is not critical. One suggestion would be to elicit an error and halt the program when the `input_signature` nesting level is found to be invalid, so the program does not progress after encountering an invalid input.

A more involved fix might include restructuring the input signature format to allow for arbitrary input nesting depth. A suggestion on this approach would be to refactor the `flattened_inputs` and `collection_inputs_` (featured in the snippet below) to be:

- A vector `std::vector<torch::jit::IValue>`, containing an ordered list of input objects to the model
- A map from `torch::jit::IValue` to `std::vector<torch_tensorrt::core::ir::Input>`, where each `IValue` is a single input to the model, and the vector stores a flattened list of Inputs

https://github.com/pytorch/TensorRT/blob/5fa63743fc4448e38b1e6b01e7403c0d74717158/core/ir/GraphInputs.cpp#L64-L66

Resolving patch:
diff --git a/py/torch_tensorrt/ts/_compile_spec.py b/py/torch_tensorrt/ts/_compile_spec.py
--- a/py/torch_tensorrt/ts/_compile_spec.py
+++ b/py/torch_tensorrt/ts/_compile_spec.py
@@ -194,17 +194,22 @@ def _parse_torch_fallback(fallback_info: Dict[str, Any]) -> _ts_C.TorchFallback:
     return info
 
 
-def _parse_input_signature(input_signature: Any):
+def _parse_input_signature(input_signature: Any, depth: int = 0):
+    if depth > 2:
+        raise AssertionError(
+            "Input nesting depth exceeds max supported depth, use 1 level: [A, B], or 2 level: [A, (B, C)]"
+        )
+
     if isinstance(input_signature, tuple):
         input_list = []
         for item in input_signature:
-            input = _parse_input_signature(item)
+            input = _parse_input_signature(item, depth + 1)
             input_list.append(input)
         return tuple(input_list)
     elif isinstance(input_signature, list):
         input_list = []
         for item in input_signature:
-            input = _parse_input_signature(item)
+            input = _parse_input_signature(item, depth + 1)
             input_list.append(input)
         return input_list
     elif isinstance(input_signature, Input) or isinstance(


Test patch: