code: def argmax_grad(x):
      argmax = tf.argmax(x)
  
      def grad(dy):
          """Let gradients pass through."""
          return tf.ones(tf.shape(x)) * dy
  
      return argmax, grad
  
imports: import tensorflow.compat.v1
label_desc: Use image node_id 3 for calling a global function with example usage: argmax_grad(x) and returns: argmax,  grad, unknown