fix(runtime): give the TensorRT engine class an __eq__ - #4584
Open
shoumikhin wants to merge 2 commits into
Open
Conversation
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.
shoumikhin
force-pushed
the
fix/engine-eq
branch
from
August 26, 2026 03:16
e1d2bdf to
dda8863
Compare
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.
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
What is broken
tests/py/dynamo/models/test_export_serde.py::test_save_load_aotiis red onmain:Why
The engine class reports
tracing_mode"real", so export hands the engineobject itself to the meta kernel instead of a fake stand-in. The engine then
ends up inside torch's fake tensor dispatch cache key. That cache hashes the
key and, on a hash match, compares with
==. The class defines no__eq__, sothe comparison raises.
Any second trace of an already compiled module hits this.
torch_tensorrt.savewith
retrace=Trueis one such path, which is why the AOT Inductor save testfails.
torch has the same problem with its own
torch::jit::OpaqueObjectand solves itthe same way, see
torch/csrc/jit/python/opaque_obj.h.Fix
Define
__eq__on the class and compare by identity. Two handles to one engineare one engine; two different engines are never equal. A comparison that says
"not equal" only costs a cache miss, so identity is both correct and the
cheapest correct answer.
Tested
This box has no TensorRT headers, so the library could not be rebuilt here.
Instead the same method was appended to the live class type at runtime, the way
torch::class_::defineMethoddoes, and the body oftest_save_load_aotiwasrun against the real engine class on an H100:
NotImplementedErrorabove, every timeA second check on a minimal custom class with the same registration shape
(
tracing_mode"real",__obj_flatten__, pickle) reproduces the failure onre-export without
__eq__and passes with it.