fixes a bug

PaddlePaddle · Feb 22, 2023 · 297d270 · 297d270
1 parent 00dccc3
commit 297d270
Show file tree

Hide file tree

Showing 2 changed files with 3 additions and 1 deletion.
diff --git a/cinn/frontend/decomposer/top_k.cc b/cinn/frontend/decomposer/top_k.cc
@@ -30,6 +30,9 @@ void top_k(const Instruction& instr, const DecomposerContext& context) {
   int k         = instr.GetAttrs<int>("k");
   CHECK_GT(k, 0) << "The attribute k must be greater than 0.";
   int axis = instr.GetAttrs<int>("axis");
+  if (axis < 0) {
+    axis += x->shape.size();
+  }
 
   auto sort_tmp    = builder->Sort(x, axis, false);
   auto sort_out    = builder->Slice(sort_tmp, {axis}, {0}, {k});

diff --git a/python/tests/ops/test_top_k_op.py b/python/tests/ops/test_top_k_op.py
@@ -75,7 +75,6 @@ def init_case(self):
 
 class TestTopKCase2(TestTopKOp):
     def init_case(self):
-        print("TestTopKCase2 ---------")
         self.inputs = {
             "x1": np.random.random([
                 2,