From dda8863615d1326d41ff1df4a9d89e53dc9ab74c Mon Sep 17 00:00:00 2001 From: Anthony Shoumikhin Date: Tue, 25 Aug 2026 19:28:08 -0700 Subject: [PATCH 1/2] fix(runtime): give the TensorRT engine class an __eq__ torch's fake tensor dispatch cache holds the engine in its key and compares keys with ==. The torchbind class had no __eq__, so the second lookup raised NotImplementedError and any re-export of a compiled module failed. Compare by identity: two handles to one engine are one engine. --- core/runtime/register_jit_hooks.cpp | 9 +++++++++ 1 file changed, 9 insertions(+) diff --git a/core/runtime/register_jit_hooks.cpp b/core/runtime/register_jit_hooks.cpp index ad334bbe37..9a4f845500 100644 --- a/core/runtime/register_jit_hooks.cpp +++ b/core/runtime/register_jit_hooks.cpp @@ -70,6 +70,15 @@ static auto TORCHTRT_UNUSED TRTEngineTSRegistrtion = .def("__str__", &TRTEngine::to_str) .def("__repr__", &TRTEngine::to_str) .def("__obj_flatten__", &TRTEngine::__obj_flatten__) + // Reporting "real" below puts the engine itself into torch's fake tensor + // dispatch cache key, and that cache compares keys with ==. Without this the + // second lookup raises "'__eq__' is not implemented", which breaks any + // re-export of a compiled module. Two handles to one engine are one engine. + .def( + "__eq__", + [](const c10::intrusive_ptr& self, const c10::intrusive_ptr& other) -> bool { + return self.get() == other.get(); + }) // Reporting "real" makes torch's tracing_with_real skip fakification and hand // the engine itself to the meta kernel, which reads only // get_serialized_metadata() -- nothing executes or mutates it. Otherwise each From b4ad44a890c1f6f8979b04feba6e288eff81664e Mon Sep 17 00:00:00 2001 From: shoumikhin Date: Tue, 25 Aug 2026 22:49:14 -0700 Subject: [PATCH 2/2] style: format test_cumsum_aten.py with black The repository lint job runs `black --check .` across the whole tree, so any file that does not match the formatter fails CI for every open pull request, not only the one that touched it. `tests/py/dynamo/conversion/test_cumsum_aten.py` is currently not black-conformant on main, which turns the Python Linting check red here. Reformat that one file with black. This is a formatting-only change: two statements that fit on a single line are un-wrapped. No test logic changes. Verified by running `black --check .` on the full tree: all files pass. --- tests/py/dynamo/conversion/test_cumsum_aten.py | 9 ++------- 1 file changed, 2 insertions(+), 7 deletions(-) diff --git a/tests/py/dynamo/conversion/test_cumsum_aten.py b/tests/py/dynamo/conversion/test_cumsum_aten.py index ab71b4339f..5c2147e281 100644 --- a/tests/py/dynamo/conversion/test_cumsum_aten.py +++ b/tests/py/dynamo/conversion/test_cumsum_aten.py @@ -33,9 +33,7 @@ def forward(self, x): ) return - self.run_test( - Cumsum(), inputs, immutable_weights=False, use_dynamo_tracer=True - ) + self.run_test(Cumsum(), inputs, immutable_weights=False, use_dynamo_tracer=True) @parameterized.expand( [ @@ -108,10 +106,7 @@ def forward(self, x): == opt_shape[positive_dim] == max_shape[positive_dim] ) - if ( - has_static_trip_count - and not is_tensorrt_rtx_version_supported("1.7") - ): + if has_static_trip_count and not is_tensorrt_rtx_version_supported("1.7"): with self.assertRaises(UnsupportedOperatorException): self.run_test_with_dynamic_shape( Cumsum(),