pytorch / pytorch/pytorch

An error should be raised when args exist in `check_inputs` but not in `inputs` for func `torch.jit.trace_module`

Open
#195,991 1 comment 0 reactions 0 assignees View on GitHub
bot-triaged oncall: jit
Dominant language
Python
Stars
103k
Forks
29.6k
PR merge metrics
PR metrics pending

Description

### 🐛 Describe the bug

The doc of `torch.jit.trace_module` is here:

https://github.com/pytorch/pytorch/blob/b45272bf18dd1448d0f7d4a351ec96c0e1f7ed90/torch/jit/_trace.py#L1043-L1137

I wonder if there should raise an error when `check_inputs` contains arg names that not in `inputs`.

That is, in the repro below, `check_inputs` contains `residual_forward` that `inputs` do not have.

I think a `KeyError` or `AttributeError` should be raised here.

### Repro
```python
import torch
import torch.nn as nn

class DynamicModule(nn.Module):
def __init__(self, dim):
super(DynamicModule, self).__init__()
self.fc1 = nn.Linear(dim, dim * 2)
self.fc2 = nn.Linear(dim * 2, dim)
self.relu = nn.ReLU()

def forward(self, x):
out = self.fc1(x)
out = self.relu(out)
out = self.fc2(out)
return out

def residual_forward(self, x):
out = self.fc1(x)
out = self.relu(out)
out = self.fc2(out)
return out + x

dim = 64
model = DynamicModule(dim)
dummy_input = torch.randn(1, dim)
traced_module = torch.jit.trace_module(
model,
inputs={'forward': dummy_input},
check_trace=True,
check_inputs=[{'forward': dummy_input * 2, 'residual_forward': dummy_input * 2}],
check_tolerance=1e-04
)
test_input = torch.randn(1, dim)
original_out = model.forward(test_input)
traced_out = traced_module.forward(test_input)
print(f"Original forward output: {original_out}")
print(f"Traced forward output: {traced_out}")
print(f"Outputs match: {torch.allclose(original_out, traced_out, atol=1e-04)}")
```
### Output
```text
Original forward output: tensor([[ 0.5123, 0.2273, -0.2485, -0.0993, -0.1016, 0.1004, 0.2824, 0.3131,
0.0802, 0.2390, -0.6841, 0.2529, -0.2849, 0.2676, 0.0185, 0.1419,
-0.0092, 0.0159, -0.2784, 0.2262, -0.5148, -0.3468, -0.4011, 0.0567,
0.4431, 0.0267, 0.2493, 0.7428, 0.1519, -0.3850, 0.2885, 0.1834,
0.5610, 0.1818, -0.0501, 0.2473, 0.3940, -0.1302, -0.1751, 0.1853,
0.0801, -0.0337, -0.0614, -0.0133, -0.4422, -0.2301, -0.0971, 0.1220,
0.1716, -0.0125, -0.0783, -0.0622, 0.2182, -0.4201, 0.2133, 0.1069,
-0.0372, 0.4152, -0.3363, -0.0935, 0.3234, 0.1637, -0.2996, -0.0625]],
grad_fn=)
Traced forward output: tensor([[ 0.5123, 0.2273, -0.2485, -0.0993, -0.1016, 0.1004, 0.2824, 0.3131,
0.0802, 0.2390, -0.6841, 0.2529, -0.2849, 0.2676, 0.0185, 0.1419,
-0.0092, 0.0159, -0.2784, 0.2262, -0.5148, -0.3468, -0.4011, 0.0567,
0.4431, 0.0267, 0.2493, 0.7428, 0.1519, -0.3850, 0.2885, 0.1834,
0.5610, 0.1818, -0.0501, 0.2473, 0.3940, -0.1302, -0.1751, 0.1853,
0.0801, -0.0337, -0.0614, -0.0133, -0.4422, -0.2301, -0.0971, 0.1220,
0.1716, -0.0125, -0.0783, -0.0622, 0.2182, -0.4201, 0.2133, 0.1069,
-0.0372, 0.4152, -0.3363, -0.0935, 0.3234, 0.1637, -0.2996, -0.0625]],
grad_fn=)
Outputs match: True
```
Thanks for noting!

### Versions

2.13.0+cu126

cc @EikanWang @jgong5 @wenzhe-nrv @sanchitintel

Contributor guide

Open the contributing guide

Research direction

The relevant entry point is torch/jit/_trace.py, in the documented trace_module section linked in the issue. Start by reproducing the supplied check_inputs/inputs mismatch and inspect how those dictionaries are validated. Done means an extra method name raises a clear error while matching method names retain existing tracing and checking behavior.

Written by the indexing model from the issue text.

Assessment

Tech stack
python
Domain
machine-learning
Issue type
Bug
Difficulty
3/5
Estimated time
1-2 days
Activity status
Active
Clarity
Mostly clear
Newbie friendliness
68/100

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.