Skip to content

fix(runtime): give the TensorRT engine class an __eq__ - #4584

Open
shoumikhin wants to merge 2 commits into
pytorch:mainfrom
shoumikhin:fix/engine-eq
Open

fix(runtime): give the TensorRT engine class an __eq__#4584
shoumikhin wants to merge 2 commits into
pytorch:mainfrom
shoumikhin:fix/engine-eq

Conversation

@shoumikhin

Copy link
Copy Markdown
Contributor

What is broken

tests/py/dynamo/models/test_export_serde.py::test_save_load_aoti is red on main:

NotImplementedError: '__eq__' is not implemented for __torch__.torch.classes.tensorrt.Engine
  File "torch/_subclasses/fake_tensor.py", line 1433, in __eq__
    return isinstance(other, _DispatchCacheKey) and self.key == other.key

Why

The engine class reports tracing_mode "real", so export hands the engine
object 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__, so
the comparison raises.

Any second trace of an already compiled module hits this. torch_tensorrt.save
with retrace=True is one such path, which is why the AOT Inductor save test
fails.

torch has the same problem with its own torch::jit::OpaqueObject and solves it
the same way, see torch/csrc/jit/python/opaque_obj.h.

Fix

Define __eq__ on the class and compare by identity. Two handles to one engine
are 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_::defineMethod does, and the body of test_save_load_aoti was
run against the real engine class on an H100:

  • without the method: the NotImplementedError above, every time
  • with the method: save, load and run succeed, cosine similarity 1.000000

A second check on a minimal custom class with the same registration shape
(tracing_mode "real", __obj_flatten__, pickle) reproduces the failure on
re-export without __eq__ and passes with it.

@meta-cla meta-cla Bot added the cla signed label Aug 26, 2026
@github-actions github-actions Bot added component: core Issues re: The core compiler component: runtime labels Aug 26, 2026
@github-actions
github-actions Bot requested a review from zewenli98 August 26, 2026 02:28
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.
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.
@github-actions github-actions Bot added the component: tests Issues re: Tests label Aug 26, 2026
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant