Skip to content

fix: Fix keep_dims functionality for aten::max #1099

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 2 commits into from
Jun 21, 2022
Merged

fix: Fix keep_dims functionality for aten::max #1099

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
84f626a
fix: Fix keep_dims functionality for aten::max
dheerajperi Jun 2, 2022
b60fe0c
chore: Fix linter issues
peri044 Jun 21, 2022
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
27 changes: 23 additions & 4 deletions core/conversion/converters/impl/max.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -18,17 +18,36 @@ auto max_registrations TORCHTRT_UNUSED = RegisterNodeConversionPatterns().patter
[](ConversionCtx* ctx, const torch::jit::Node* n, args& args) -> bool {
auto self = args[0].ITensorOrFreeze(ctx);
auto dim = args[1].unwrapToInt();
auto keep_dims = args[2].unwrapToBool();
auto selfDim = util::toVec(self->getDimensions());
if (dim < 0) {
dim = selfDim.size() + dim;
}
uint32_t shiftDim = 1 << dim;
auto TopKOperation = nvinfer1::TopKOperation::kMAX;
auto new_layer = ctx->net->addTopK(*self, TopKOperation, 1, shiftDim);
TORCHTRT_CHECK(new_layer, "Unable to create max layer from node: " << *n);
auto topk_layer = ctx->net->addTopK(*self, TopKOperation, 1, shiftDim);
TORCHTRT_CHECK(topk_layer, "Unable to create max layer from node: " << *n);
auto topk_dims = util::toVec(topk_layer->getOutput(0)->getDimensions());

auto out0 = ctx->AssociateValueAndTensor(n->outputs()[0], new_layer->getOutput(0));
auto out1 = ctx->AssociateValueAndTensor(n->outputs()[1], new_layer->getOutput(1));
nvinfer1::ITensor* out0;
nvinfer1::ITensor* out1;
if (!keep_dims) {
if (topk_dims[dim] == 1) {
auto squeeze_layer = ctx->net->addShuffle(*topk_layer->getOutput(0));
squeeze_layer->setReshapeDimensions(util::squeezeDims(topk_layer->getOutput(0)->getDimensions(), dim));
TORCHTRT_CHECK(squeeze_layer, "Unable to create squeeze_layer layer from node: " << *n);
out0 = ctx->AssociateValueAndTensor(n->outputs()[0], squeeze_layer->getOutput(0));

auto squeeze_layer_indices = ctx->net->addShuffle(*topk_layer->getOutput(1));
squeeze_layer_indices->setReshapeDimensions(
util::squeezeDims(topk_layer->getOutput(1)->getDimensions(), dim));
TORCHTRT_CHECK(squeeze_layer_indices, "Unable to create squeeze_layer_indices layer from node: " << *n);
out1 = ctx->AssociateValueAndTensor(n->outputs()[1], squeeze_layer_indices->getOutput(0));
}
} else {
out0 = ctx->AssociateValueAndTensor(n->outputs()[0], topk_layer->getOutput(0));
out1 = ctx->AssociateValueAndTensor(n->outputs()[1], topk_layer->getOutput(1));
}

LOG_DEBUG("Output tensor(0) shape: " << out0->getDimensions());
LOG_DEBUG("Output tensor(1) shape: " << out1->getDimensions());
Expand Down
12 changes: 12 additions & 0 deletions tests/core/conversion/converters/test_reduce.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -212,6 +212,18 @@ TEST(Converters, ATenProdKeepDimsConvertsCorrectly) {
test_body(graph, in);
}

TEST(Converters, ATenMaxKeepDimsConvertsCorrectly) {
const auto graph = R"IR(
graph(%x : Tensor):
%2 : int = prim::Constant[value=-1]()
%3 : bool = prim::Constant[value=1]()
%keep.1 : Tensor, %6 : Tensor = aten::max(%x, %2, %3)
return (%keep.1, %6))IR";

auto in = at::randint(-5, 5, {4, 4}, at::kCUDA);
test_body(graph, in);
}

TEST(Converters, ATenMeanDimNegOneIndexConvertsCorrectly) {
const auto graph = R"IR(
graph(%0 : Tensor):
Expand Down