tf.argmax函数

`tf.argmax` 是 TensorFlow 中的一个函数，它用于获取指定轴上元素最大值的索引。这个函数可以帮助你在张量（tensor）中快速找到每行（或每列，或其他维度）的最大值所在的索引位置。

`tf.argmax` 的基本用法如下：
- `tf.argmax(input, axis=None, output_type=tf.int64, name=None)`
- `input` 是输入的张量。
- `axis` 是张量的轴，用于指定在哪个轴上进行求索引操作。默认值为 None，此时会返回扁平化张量中的最大值索引。
- `output_type` 指定输出索引的数据类型，默认为 `tf.int64`。
- `name` 是可选的参数，用于给操作命名。

例如，如果你有一个形状为 [3, 3] 的二维张量 `matrix`，并且想要获取每一行的最大值索引，你可以使用 `tf.argmax` 函数并设置 `axis=1`。

使用时的示例代码如下：

import tensorflow as tf

# 创建一个3x3的矩阵
matrix = tf.constant([
    [1, 5, 3],
    [2, 8, 6],
    [7, 4, 9]
])

# 使用tf.argmax获取每行最大值的索引
max_indices = tf.argmax(matrix, axis=1)

# 会得到一个包含每行最大值索引的向量[1, 1, 2]