code: def test_topk_backward(device_id, precision):
      def check_grad_last_axis(input, root, indices, output):
          d = input.shape[-1]
          k = indices.shape[-1]
          expected_output = np.zeros_like(input).reshape(-1, d)
          ind = np.reshape(indices, (-1, k))
          r = np.reshape(root, (-1, k))
          assert ind.shape[0] == r.shape[0] == expected_output.shape[0]
          for i in range(expected_output.shape[0]):
              for j in range(k):
                  expected_output[i, int(ind[i, j])] = r[i, j]
          expected_output = expected_output.reshape(input.shape)
          assert np.allclose(output, expected_output)
  
      dt = PRECISION_TO_TYPE[precision]
      dev = cntk_device(device_id)
  
      axis = -1
      h = C.placeholder()
      p = C.parameter((4, 5, 6))
      p.value = p.value + np.random.randn(*p.shape)
      y = C.top_k(h, 3, axis=axis)
      y.replace_placeholder(p)
      dy, top = y.forward({}, y.outputs, set([y.outputs[0]]))
      indices = top[y.outputs[1]]
      root = np.ones_like(indices)
      root = root + np.arange(np.prod(root.shape)).reshape(*root.shape)
      cg = y.backward(dy, {y.outputs[0]: root}, set([p]))[p]
      check_grad_last_axis(p.value, root, indices, cg)
  
      q = C.sequence.input_variable((5, 6), needs_gradient=True)
      q0 = [np.random.randn(4 - i, 5, 6).astype(dt) for i in range(2)]
      y = C.top_k(q, 3, axis=axis)
      dy, top = y.forward(
          {q: q0}, y.outputs, set([y.outputs[0]]), device=dev
      )
      indices = top[y.outputs[1]]
      root = [
          np.ones_like(i)
          + 100 * k
          + np.arange(np.prod(i.shape)).reshape(*i.shape)
          for k, i in enumerate(indices)
      ]
      cg = y.backward(dy, {y.outputs[0]: root}, set([q]))[q]
      for i in range(2):
          check_grad_last_axis(q0[i], root[i], indices[i], cg[i])
  
imports: from __future__ import division
  import numpy
  import cntk
  import pytest
  from .ops_test_utils import unittest_helper, _test_unary_op, _test_binary_op, AA, precision, PRECISION_TO_TYPE, cntk_device
  import cntk
  from cntk import Value
  from cntk.axis import Axis
  from cntk.internal import sanitize_dtype_cntk
  from ..None import constant
label_desc: Use image node_id 21 for calling a global function with example usage: test_topk_backward(device_id, precision) without return types