fix: run select with a computed index - #4709
Open
shoumikhin wants to merge 1 commit into
Open
shoumikhin wants to merge 1 commit into
shoumikhin wants to merge 1 commit into
Conversation
github-actions
Bot
added
component: tests
## Problem `select` cannot run in TensorRT when its index is computed from another input. The index also needs the right shape: selecting one row from an `(8, 4)` tensor should return `(4,)`, not `(1, 4)`. ## Change Accept an index computed at runtime. Convert a one-element index tensor to a scalar (a value with no tensor dimensions) before passing it to TensorRT's gather operation. Record how many dimensions each scalar input expects. Use that information when calling the engine and when exporting the compiled model. This keeps integer and floating-point scalar inputs consistent across those paths. Check that the model test builds an engine of the expected type. Runtime selection depends on native-library availability; the deprecated `use_python_runtime` option does not select a backend. ## Tests Passed 7/7 converter tests. The computed-index model and a floating-point scalar control both passed in separate native and Python-only package layouts. Assertions checked native-library availability and confirmed native `ScriptObject` engines or Python `TRTEngine` engines. The compiled wrapper ran directly with inputs producing indices 3, 5, and 0. Shape, data type, and values matched ordinary PyTorch. A separate export check also ran the wrapper directly. It then used both exporters, legacy and nonlegacy, to export, save, reload, and run the model with the same changing indices. All four exporter/backend combinations passed. Saving used `BytesIO`, an in-memory byte stream, not a file on disk. That export check explicitly limits indices to the valid range. It does not test out-of-range behavior. Existing gather bounds behavior and conversion of indices to 32-bit integers are unchanged. Tested on Linux x86_64 with Python 3.12, PyTorch 2.15 nightly, and TensorRT 11.2. Windows, aarch64, TensorRT-RTX, TensorRT 10.x, changing input dimensions, inputs on different devices, concurrent calls, older saved models, and filesystem durability were not tested.
shoumikhin
force-pushed
the
upstream/select-runtime-index
branch
from
September 12, 2026 04:58
a854f6f to
a54a73e
Compare
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
1 participant
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.
Problem
selectcannot run in TensorRT when its index is computed from another input. The index also needs the right shape: selecting one row from an(8, 4)tensor should return(4,), not(1, 4).Change
Accept an index computed at runtime. Convert a one-element index tensor to a scalar before passing it to TensorRT's gather operation.
Record how many dimensions each scalar input expects. Use that information when calling the engine and when exporting the compiled model. This keeps integer and floating-point scalar inputs consistent across those paths.
Check that the model test builds an engine of the expected type. Runtime selection depends on native-library availability; the deprecated
use_python_runtimeoption does not select a backend.Tests
Passed 7/7 converter tests with each runtime. The computed-index model passed with both; a floating-point scalar control also passed with the native runtime. Earlier testing covered that scalar control with the Python runtime too. Assertions confirmed native
ScriptObjector PythonTRTEngineengines. The converter control fails on the base branch when it tries to convert a TensorRT index into a NumPy constant.Direct execution matched PyTorch shape, data type, and values for indices 3, 5, and 0. A separate check used both exporters, legacy and nonlegacy, to export, save to
BytesIO, reload, and execute the same changing indices. All four exporter/runtime combinations passed.The rerun used Linux x86_64, Python 3.12, and PyTorch 2.15 nightly. Python-runtime checks used TensorRT 11.3. Native checks used the published Torch-TensorRT 2.15 nightly library with its TensorRT 11.2 dependency; they do not validate a native TensorRT 11.3 build.
Export checks bound indices to the valid range. Existing gather bounds behavior and conversion to 32-bit indices are unchanged. Windows, aarch64, TensorRT-RTX, TensorRT 10.x, changing input dimensions, different input devices, concurrent calls, older saved models, and filesystem durability were not tested.