[onnx] Fix crash when negative transpose values exist (#2915)

We are crashing due to indexing into a negative shape. Updated the
lowering to avoid the crash.
pull/2922/head
Rob Suderman 2024-02-16 13:04:47 -08:00 committed by GitHub
parent c5d8c12469
commit 468c533942
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
1 changed files with 11 additions and 1 deletions

View File

@ -1447,7 +1447,17 @@ void mlir::torch::onnx_c::populateDefaultDomainQtoZ(
current[i] = i;
}
// Convert dynamic shape dimension.
for (auto &dim : permutations)
dim = dim < 0 ? dim + rank : dim;
// We need to override to the destination if known:
if (resultType.hasSizes()) {
for (int i = 0; i < rank; ++i) {
shape[permutations[i]] = resultType.getSizes()[i];
}
}
// Convert dynamic shape dimension:
for (unsigned i = 0; i < shape.size(); i++) {
if (shape[i] == ShapedType::kDynamic)
shape[i] = Torch::kUnknownSize;