如何在TensorFlow中编写argmax函数？

Question

I'm not talking about tf.argmax but about argmax in the mathematical sense, eg given a discrete set of values and a function find the value that maximizes it. I currently have something like this

input = tf.placeholder(name='input')
Qhat = do_stuff_to(input) # e.g. tf.add(input, 3)

Now I want to define another TensorFlow node, max_Qhat , which will take an array of Tensors as its argument. It will feed each of these tensors to Qhat and return the one that resulted in the greatest value. How do I do this? ( NOTE I am not looking to run Qhat, so no session.run . I just want to define a function that evaluates it.)

Code I have so far:

inputs = tf.placeholder(name='inputs')
max_Qhat = ???

Answer 1

Create an input tensor, say of dimensions [k, ...] (this is your "array" of k input tensors that we "maximize" over) then compute the Qhat op over the first axis/dimension of the input tensor (eg. using tf.map_fn ) such that it returns a tensor of function values of dimensions [k] and then determine the maximum of this returned tensor.

import tensorflow as tf;

sess = tf.InteractiveSession();

# define inputs
inputs = tf.constant([
    [1, 2, 3, 4, 5],
    [4, 3, 6, 2, 1],
    [9, 9, 9, 9, 9],
    [0, 1, 3, 5, 2]
], shape = (4, 5));

# the function/op that we want to compute for each input tensor
def some_op(t):
    return tf.reduce_sum(t);

# compute the function values
q = tf.map_fn(some_op, inputs);

# determine the index of the input tensor that maximizes the function
index = sess.run(tf.argmax(q, axis = 0));
maximizer = sess.run(inputs[index]);