已合并
Adapt test_jit_is_trancing in trace for NPU #37080
yrouyi创建于 5月29日
Adapt test_jit_is_trancing in trace for NPU #37080
已合并
共 1 个文件变更+52-0
| @@ -0,0 +1,52 @@ | |||
| 1 | +# Owner(s): ["oncall: jit"] | ||
| 2 | +""" | ||
| 3 | +Add validation cases for torch.jit tracing APIs on NPU: | ||
| 4 | +Strictly validates the functional correctness of torch.jit.is_tracing() | ||
| 5 | +during torch.jit.trace execution. | ||
| 6 | +Note: Script mode already has community test cases in test/test_jit.py, | ||
| 7 | +this case only verifies trace mode. | ||
| 8 | +""" | ||
| 9 | + | ||
| 10 | +import torch | ||
| 11 | +from torch.testing._internal.common_utils import run_tests, TestCase | ||
| 12 | +import torch_npu | ||
| 13 | + | ||
| 14 | +# Get current accelerator device type | ||
| 15 | +device_type = acc.type if (acc := torch.accelerator.current_accelerator()) else "cpu" | ||
| 16 | + | ||
| 17 | + | ||
| 18 | +class TestJitIsTracing(TestCase): | ||
| 19 | + def test_is_tracing_returns_true_in_trace_mode(self): | ||
| 20 | + """ | ||
| 21 | + Validates that torch.jit.is_tracing() returns True during trace recording. | ||
| 22 | + | ||
| 23 | + Note: Direct assertion inside the traced function is not feasible because | ||
| 24 | + the Python-level tracing flag remains False during Eager-mode execution | ||
| 25 | + within torch.jit.trace. Instead, this test strictly verifies the API's | ||
| 26 | + correctness through behavioral validation: | ||
| 27 | + | ||
| 28 | + If is_tracing() correctly evaluates to True during recording, the tracer | ||
| 29 | + will capture the 'x + 1' branch and permanently bake it into the TorchScript graph. | ||
| 30 | + We assert the final output to prove this specific path was recorded. | ||
| 31 | + """ | ||
| 32 | + def my_func(x): | ||
| 33 | + # The tracer evaluates this condition during the recording phase. | ||
| 34 | + if torch.jit.is_tracing(): | ||
| 35 | + return x + 1 | ||
| 36 | + else: | ||
| 37 | + return x - 1 | ||
| 38 | + | ||
| 39 | + inp = torch.randn(3, 3).to(device_type) | ||
| 40 | + | ||
| 41 | + # Execute tracing. If is_tracing() works correctly, the 'x + 1' path is recorded. | ||
| 42 | + traced_func = torch.jit.trace(my_func, inp, check_trace=False) | ||
| 43 | + | ||
| 44 | + # Strict behavioral proof: The output MUST follow the 'x + 1' branch, | ||
| 45 | + # proving that is_tracing() returned True and the path was successfully captured. | ||
| 46 | + traced_output = traced_func(inp) | ||
| 47 | + self.assertEqual(traced_output, inp + 1, | ||
| 48 | + msg="Traced model did not record the is_tracing()==True branch.") | ||
| 49 | + | ||
| 50 | + | ||
| 51 | +if __name__ == "__main__": | ||
| 52 | + run_tests() | ||