Skip to content

[fix] aten::stack with dynamic inputs #1804

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
merged 1 commit into from
Apr 7, 2023
Merged

[fix] aten::stack with dynamic inputs #1804

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
3 changes: 1 addition & 2 deletions core/conversion/converters/impl/stack.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -43,10 +43,9 @@ auto stack_registrations TORCHTRT_UNUSED = RegisterNodeConversionPatterns().patt
auto cont = t.toCustomClass<TensorContainer>();
itensor = cont->tensor();
}

auto shuffle_layer = ctx->net->addShuffle(*itensor);
TORCHTRT_CHECK(shuffle_layer, "Unable to create shuffle layer from node: " << *n);
shuffle_layer->setReshapeDimensions(util::unsqueezeDims(itensor->getDimensions(), dim));
shuffle_layer->setReshapeDimensions(util::unsqueezeDims(itensor->getDimensions(), dim, 1, false));
Copy link
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

In the description of this PR, it says Unsqueezed dims used 0 for dynamic dimension which was not aligned correctly with the dynamic dim in the input, switch to using -1.. Here the val is 1 (which is typical of an unsqueeze op). Can you elaborate on what's going on here related to dynamic shapes?

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I'm referring to the last arg to unsqueezeDims use_zeros which previously defaulted to true and converted dynamic dimensions to 0 rather than -1.


tensors.push_back(shuffle_layer->getOutput(0));
}
Expand Down
35 changes: 35 additions & 0 deletions tests/core/conversion/converters/test_stack.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -40,6 +40,41 @@ TEST(Converters, ATenStackPureTensorConvertsCorrectly) {
TestATenStackPureTensorConvertsCorrectly(graph2);
}

TEST(Converters, ATenStackPureTensorDynamicConvertsCorrectly) {
auto TestATenStackPureTensorConvertsCorrectly = [](const std::string& graph) {
auto g = std::make_shared<torch::jit::Graph>();
torch::jit::parseIR(graph, g.get());

auto in1 = at::randint(1, 10, {4, 4, 4}, {at::kCUDA});
auto in2 = at::randint(1, 10, {4, 4, 4}, {at::kCUDA});

auto params = torch_tensorrt::core::ir::get_static_params(g->inputs(), {});
auto jit_results = torch_tensorrt::tests::util::RunGraph(g, params, {in1, in2});

params = torch_tensorrt::core::ir::get_static_params(g->inputs(), {});
auto trt_results = torch_tensorrt::tests::util::RunGraphEngineDynamic(g, params, {in1, in2});

ASSERT_TRUE(torch_tensorrt::tests::util::almostEqual(jit_results[0], trt_results[0], THRESHOLD_E5));
};
const auto graph = R"IR(
graph(%0 : Tensor,
%1 : Tensor):
%2 : Tensor[] = prim::ListConstruct(%0, %1)
%3 : int = prim::Constant[value=1]()
%4 : Tensor = aten::stack(%2, %3)
return (%4))IR";
const auto graph2 = R"IR(
graph(%0 : Tensor,
%1 : Tensor):
%2 : Tensor[] = prim::ListConstruct(%0, %1)
%3 : int = prim::Constant[value=-1]()
%4 : Tensor = aten::stack(%2, %3)
return (%4))IR";

TestATenStackPureTensorConvertsCorrectly(graph);
TestATenStackPureTensorConvertsCorrectly(graph2);
}

TEST(Converters, ATenStackDiffTensorConvertsCorrectly) {
auto TestATenStackDiffTensorConvertsCorrectly = [](const std::string& graph) {
auto g = std::make_shared<torch::jit::Graph>();
Expand Down