TensorFlow tf.argmax() does not Support Backprop and Gradient Operation – TensorFlow Tutorial

TensorFlow tf.argmax() can allow us to get the index with the largest value across axes of a tensor, which is widely used in classification problems. However, does it support backprop and gradient operation in tensorflow? We will discuss this topic with an example in this tutorial.

Please look at code below.

Calculate z = x * y

First, we will create two tensor x and y, then get z by z = x * y. Here is an example:

import tensorflow as tf
import numpy as np

x = tf.Variable(np.array([[1, 9, 3],[4, 5, 6]]), dtype = tf.float32)
y = tf.Variable(np.array([[1, 1],[5, 2], [2, 7]]), dtype = tf.float32)

z = tf.matmul(x, y)

Calculate softmax value of z on axis = 1

z = tf.nn.softmax(z, axis = 1)

Get the prediction value based on z

predict = tf.argmax(z, axis = 1)

where predict is the probable labels in classification problems.

Get the m tensor by predict

c = tf.Variable(np.array([[2, 1],[5, 5], [2, 2]]), dtype = tf.float32)

m = tf.nn.embedding_lookup(c, ids=predict)

Calculate u by u = m * n

Here is example code.

n = tf.Variable(np.array([[1, 2],[2, 2]]), dtype = tf.float32)

u = tf.matmul(m,n)

Compute the gradient of x based on u

r = tf.gradients(u, x)

Output gradient

init = tf.global_variables_initializer() 
init_local = tf.local_variables_initializer()

with tf.Session() as sess:
    sess.run([init, init_local])
    print(sess.run([r]))

We have known tf.nn.embedding_lookup() support gradient operation first.

Run this code, we will get result: TypeError: Fetch argument None has invalid type <class ‘NoneType’>

It means tensorflow tf.argmax() function does not support support backprop and bradient operation. You can not use it in compute flow of models.