Python 如何在TensorFlow中编写argmax函数？

我说的不是

tf.argmax

，而是数学意义上的argmax，例如，给定一组离散值和一个函数，找到使其最大化的值。我现在有类似的东西

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

现在我想定义另一个TensorFlow节点，

max_Qhat

，它将使用一个张量数组作为参数。它将这些张量中的每一个馈送到

Qhat

，并返回产生最大值的张量。我该怎么做？（注意我不想运行Qhat，所以没有

会话。运行

。我只想定义一个函数来计算它。）

到目前为止，我掌握的代码是：

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

---

创建一个输入张量，比如维度

[k，…]

（这是我们“最大化”的

k

输入张量的“数组”），然后计算输入张量的第一个轴/维度上的Qhat op（例如，使用

tf.map\u fn

），以便返回维度的函数值的张量

[k]

然后确定此返回张量的最大值

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]);