TensorFlow: Макс тензора вдоль оси

Мой вопрос состоит из двух связанных частей:

Как рассчитать максимум вдоль некоторой оси тензора? Например, если у меня есть
```
x = tf.constant([[1,220,55],[4,3,-1]])
```
Мне нужно что-то вроде
```
x_max = tf.max(x, axis=1)
print sess.run(x_max)

output: [220,4]
```
Я знаю, что есть tf.argmax и a tf.maximum, но не дают максимального значения вдоль оси одного тензора. На данный момент у меня есть обходное решение:
```
x_max = tf.slice(x, begin=[0,0], size=[-1,1])
for a in range(1,2):
    x_max = tf.maximum(x_max , tf.slice(x, begin=[0,a], size=[-1,1]))
```
Но он выглядит менее оптимальным. Есть ли лучший способ сделать это?
Учитывая индексы тензора a argmax, как я могу индексировать в другой тензор, используя эти индексы? Используя пример x выше, как мне сделать что-то вроде следующего:
```
ind_max = tf.argmax(x, dimension=1)    #output is [1,0]
y = tf.constant([[1,2,3], [6,5,4])
y_ = y[:, ind_max]                     #y_ should be [2,6]
```
Я знаю, что нарезка, как и последняя строка, еще не существует в TensorFlow (# 206).

Мой вопрос: каково наилучшее обходное решение для моего конкретного случая (возможно, используя другие методы, такие как сбор, выбор и т.д.)?

Дополнительная информация: Я знаю, что x и y будут только двухмерными тензорами!

Ответ 1

Оператор tf.reduce_max() обеспечивает именно эту функциональность. По умолчанию он вычисляет глобальный максимум данного тензора, но вы можете указать список reduction_indices, который имеет то же значение, что и axis в NumPy. Чтобы завершить свой пример:

x = tf.constant([[1, 220, 55], [4, 3, -1]])
x_max = tf.reduce_max(x, reduction_indices=[1])
print sess.run(x_max)  # ==> "array([220,   4], dtype=int32)"

Если вы вычисляете argmax с помощью tf.argmax(), вы можете получить значения из другого тензора y, сглаживая y, используя tf.reshape(), преобразуя индексы argmax в векторные индексы следующим образом и используя tf.gather(), чтобы извлечь соответствующие значения:

ind_max = tf.argmax(x, dimension=1)
y = tf.constant([[1, 2, 3], [6, 5, 4]])

flat_y = tf.reshape(y, [-1])  # Reshape to a vector.

# N.B. Handles 2-D case only.
flat_ind_max = ind_max + tf.cast(tf.range(tf.shape(y)[0]) * tf.shape(y)[1], tf.int64)

y_ = tf.gather(flat_y, flat_ind_max)

print sess.run(y_) # ==> "array([2, 6], dtype=int32)"