code: def _testArg(
      self,
      method,
      x,
      axis,
      expected_values,
      use_gpu=False,
      expected_err_re=None,
  ):
      with self.session(use_gpu=use_gpu):
          ans = method(x, axis=axis)
          if expected_err_re is None:
              tf_ans = self.evaluate(ans)
              # Defaults to int64 output.
              self.assertEqual(np.int64, tf_ans.dtype)
              self.assertAllEqual(tf_ans, expected_values)
              self.assertShapeEqual(expected_values, ans)
          else:
              with self.assertRaisesOpError(expected_err_re):
                  self.evaluate(ans)
  
imports: import functools
  from absl.testing import parameterized
  import numpy
  from tensorflow.python.framework import dtypes
  from tensorflow.python.framework import test_util
  from tensorflow.python.ops import array_ops
  from tensorflow.python.ops import math_ops
  from tensorflow.python.platform import test
label_desc: Use image node_id 1 for calling the ArgMaxTest obj's underlying member method code with example usage: obj._testArg(method, x, axis, expected_values, use_gpu, expected_err_re) without return types