From e5d46d8db810708620c22d12eb8e505b31b98a24 Mon Sep 17 00:00:00 2001 From: gs-olive <113141689+gs-olive@users.noreply.github.com> Date: Wed, 14 Jun 2023 14:37:58 -0700 Subject: [PATCH] fix: Repair flaky TopK core test - Update test code to generate a random unique sequence of integers for TopK comparison testing, avoiding index and value collisions in the results which were causing CI test failures --- tests/core/conversion/converters/test_topk.cpp | 13 ++++++++++++- 1 file changed, 12 insertions(+), 1 deletion(-) diff --git a/tests/core/conversion/converters/test_topk.cpp b/tests/core/conversion/converters/test_topk.cpp index 9117fe1bb2..c1f651b3bc 100644 --- a/tests/core/conversion/converters/test_topk.cpp +++ b/tests/core/conversion/converters/test_topk.cpp @@ -17,7 +17,18 @@ TEST(Converters, ATenTopKConvertsCorrectly) { auto g = std::make_shared(); torch::jit::parseIR(graph, g.get()); - auto in = at::rand({10, 10, 100}, {at::kCUDA}); + auto dim0 = 10, dim1 = 10, dim2 = 100; + + // Initialize zero tensor to be filled with random indices along the final dimension + auto in = at::zeros({dim0, dim1, dim2}, {at::kCUDA}); + + // For each final dimension, fill it with random scramble of unique integers in the range [0, dim0*dim1*dim2) + for (auto i = 0; i < dim0; i++) { + for (auto j = 0; j < dim1; j++) { + auto random_index_permutation = at::randperm(dim0 * dim1 * dim2, c10::kInt, {}, at::kCUDA, {}).slice(0, 0, dim2); + in.slice(0, i, i + 1).slice(1, j, j + 1) = random_index_permutation; + } + } auto params = torch_tensorrt::core::ir::get_static_params(g->inputs(), {}); auto jit_results = torch_tensorrt::tests::util::RunGraph(g, params, {in});