text: #  Copyright (c) 2021 PaddlePaddle Authors. All Rights Reserved.
  #
  # Licensed under the Apache License, Version 2.0 (the "License");
  # you may not use this file except in compliance with the License.
  # You may obtain a copy of the License at
  #
  #    http://www.apache.org/licenses/LICENSE-2.0
  #
  # Unless required by applicable law or agreed to in writing, software
  # distributed under the License is distributed on an "AS IS" BASIS,
  # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
  # See the License for the specific language governing permissions and
  # limitations under the License.
  
  import unittest
  
  import numpy as np
  import paddle
  import paddle.fluid as fluid
  import paddle.fluid.compiler as compiler
  import paddle.optimizer
  import paddle.static
  from paddle.fluid.tests.unittests.ipu.op_test_ipu import IPUOpTest
  
  paddle.enable_static()
  
  
  @unittest.skipIf(not paddle.is_compiled_with_ipu(),
                   "core is not compiled with IPU")
  class TestTopKOp(IPUOpTest):
      def setUp(self):
          self.set_ops()
          self.set_atol()
          self.set_training()
          self.k = 3
          self.use_K_as_const_variable = False
  
          self.set_feed()
          self.set_attrs()
  
      def set_ops(self):
          self.ops = [
              paddle.fluid.layers.topk,
              paddle.topk  # use top_k_v2 implementation
          ]
  
      def set_feed(self):
          self.feed_shape = []
          self.feed_shape.append([3, 5])
  
          self.feed = {}
          self.feed_list = []
          self.feed["in_0"] = np.random.uniform(
              size=self.feed_shape[0]).astype(np.float32)
          self.feed_list.append("in_0")
          if self.use_K_as_const_variable:
              # self.feed["in_1"] = np.array([self.k]).astype("int32")
              # self.feed_list.append("in_1")
              pass
          print("[TestTopKop] feed data:\n%s" % self.feed["in_0"])
  
      def set_attrs(self):
          self.attrs = {
              # "axis": -1,
              # "sorted": True
          }
          if not self.use_K_as_const_variable:
              self.attrs["k"] = self.k
  
      def _test_base(self, run_ipu=True, op=None, data_feed=None):
          assert (op is not None)
          assert (data_feed is not None)
          scope = fluid.core.Scope()
          main_prog = paddle.static.Program()
          startup_prog = paddle.static.Program()
          SEED = self.SEED
          main_prog.random_seed = SEED
          startup_prog.random_seed = SEED
  
          with fluid.scope_guard(scope):
              with paddle.static.program_guard(main_prog, startup_prog):
                  x = paddle.static.data(
                      name=self.feed_list[0],
                      shape=self.feed_shape[0],
                      dtype='float32')
                  if not self.use_K_as_const_variable:
                      topk_values, topk_indices = op(x, **self.attrs)
                  else:
                      # !important, popart cannot accept non const tensor
                      # K_t = paddle.static.data(name="in_1", shape=[1], dtype='int32')
                      K_t = fluid.layers.fill_constant(
                          shape=[1], dtype='int32', value=self.k, name="in_2")
                      topk_values, topk_indices = op(x, K_t, **self.attrs)
                  fetch_list = [topk_values.name, topk_indices.name]
  
              if run_ipu:
                  place = paddle.IPUPlace()
              else:
                  place = paddle.CPUPlace()
              exe = paddle.static.Executor(place)
              exe.run(startup_prog)
  
              if run_ipu:
                  feed_list = self.feed_list
                  ipu_strategy = paddle.static.IpuStrategy()
                  ipu_strategy.SetGraphConfig(is_training=self.is_training)
                  program = compiler.IPUCompiledProgram(
                      main_prog,
                      ipu_strategy=ipu_strategy).compile(feed_list, fetch_list)
              else:
                  program = main_prog
  
              print("Running inference ...")
              result = exe.run(program, feed=data_feed, fetch_list=fetch_list)
              print("Complete running infrence.")
              return result
  
      def test_base(self):
          for op in self.ops:
              res0_topk_values, res0_topk_indices = self._test_base(
                  True, op=op, data_feed=self.feed)
              res1_topk_values, res1_topk_indices = self._test_base(
                  False, op=paddle.fluid.layers.topk, data_feed=self.feed)
  
              print("[TestTopKop] IPU res0 values:\n%s\n" % res0_topk_values)
              print("[TestTopKop] CPU res1 values:\n%s\n" % res1_topk_values)
              view_type = np.uint32
              print("[TestTopKop] IPU res0 indices:\n%s\n" %
                    res0_topk_indices.astype(view_type))
              print("[TestTopKop] CPU res1 indices:\n%s\n" % res1_topk_indices)
  
              self.assertTrue(
                  np.allclose(
                      res0_topk_values.flatten(),
                      res1_topk_values.flatten(),
                      atol=self.atol))
  
              self.assertTrue(
                  np.allclose(
                      res0_topk_indices.astype(view_type).flatten(),
                      res1_topk_indices.flatten(),
                      atol=self.atol))
  
  
  if __name__ == "__main__":
      unittest.main()
  