code: def test_argtopk_nd(self, shp, k_, dtype, sorted, idx_dtype):
      ndim = len(shp)
      for axis in range(-ndim, ndim):
          if isinstance(k_, str):
              k = eval(k_.replace("n", str(shp[axis])))
          else:
              k = k_
  
          if k == 0:
              continue
  
          x = theano.tensor.tensor(
              name="x", broadcastable=(False,) * len(shp), dtype=dtype
          )
          y = argtopk(
              x, k, axis=axis, sorted=sorted, idx_dtype=idx_dtype
          )
          fn = theano.function([x], y, mode=self.mode)
          assert any(
              [
                  isinstance(n.op, self.op_class)
                  for n in fn.maker.fgraph.apply_nodes
              ]
          )
          size = reduce(int.__mul__, shp)
          xval = gen_unique_vector(size, dtype).reshape(shp)
          yval = fn(xval)
          idx = slice(-k, None) if k > 0 else slice(-k)
          l = axis % ndim
          r = ndim - l
          idx = (slice(None),) * l + (idx,) + (slice(None),) * (r - 1)
          goal = np.argsort(xval, axis=axis)[idx].astype(idx_dtype)
  
          assert np.all(
              np.sort(yval, axis=axis) == np.sort(goal, axis=axis)
          )
  
imports: from __future__ import absolute_import, print_function, division
  from itertools import product, chain
  from functools import reduce
  import unittest
  from theano.tests import unittest_tools
  import numpy
  import theano
  from theano import tensor
  from theano.tensor.sort import sort, SortOp
  from theano.tensor.sort import argsort, ArgSortOp
  from theano.tensor.sort import topk, argtopk, topk_and_argtopk, TopKOp
label_desc: Use image node_id 9 for calling the Test_TopK obj's underlying member method code with example usage: obj.test_argtopk_nd(shp, k_, dtype, sorted, idx_dtype) without return types