code: def test_top_k_grad(self, X, k, axis, gc, dc):
      dims = X.shape
      if axis >= len(dims):
          axis %= len(dims)
  
      input_axis = len(dims) - 1 if axis == -1 else axis
      prev_dims = 1
      next_dims = 1
      for i in range(input_axis):
          prev_dims *= dims[i]
      for i in range(input_axis + 1, len(dims)):
          next_dims *= dims[i]
  
      X_flat = X.reshape((prev_dims, dims[input_axis], next_dims))
      for i in range(prev_dims):
          for j in range(next_dims):
              # this try to make sure adding stepsize (0.05)
              # will not change TopK selections at all
              X_flat[i, :, j] = (
                  np.arange(dims[axis], dtype=np.float32) / 5
              )
              np.random.shuffle(X_flat[i, :, j])
      X = X_flat.reshape(dims)
  
      op = core.CreateOperator(
          "TopK",
          ["X"],
          ["Values", "Indices"],
          k=k,
          axis=axis,
          device_option=gc,
      )
  
      self.assertGradientChecks(gc, op, [X], 0, [0], stepsize=0.05)
  
file_path: pytorch/caffe2/python/operator_test/top_k_test.py
imports: import hypothesis.strategies
  import numpy
  from caffe2.python import core
  from hypothesis import given, settings
  import caffe2.python.hypothesis_test_util
  import caffe2.python.serialized_test.serialized_test_util
label_desc: Use image node_id 10 for calling the TestTopK obj's underlying member method code with example usage: obj.test_top_k_grad(X, k, axis, gc, dc) without return types