TensorFlow 1.14を使用して3Dテンソルから最大の2Dテンソルを取得するにはどうすればよいですか？

Question

TensorFlow 1.14を使用して、最大1つの値に基づいて3Dテンソルから2D最大テンソルを取得するための最良かつ最適化された方法（ループなし）を探しています。このテンソルとこの関数があるとしましょう（理解のために-それは機能していません-）：

def get_Max(inputs):
    max_indices = [0,0,0]
    for i in range(16):
        for j in range(2048):
            for k in range(10):
                if(inputs[max_indices[0],max_indices[1],max_indices[2]]<inputs[i,j,k]):
                   max_indices = [i,j,k]
    return inputs[:][j]
inputs = tf.random.uniform(shape=[16,2048,10],dtype=tf.dtypes.float32)
output = get_Max(inputs)

したがって、出力テンソルは2048からの16の最大値である[16,10]の形状である必要があります。では、ループなしでこれを実行できる関数をどのように実装できますか？

私は使用しましたtf.math.reduce_maxが、下の画像で明らかなように、それは私が探しているものではありません：

Answer 1

inp = tf.random.uniform(shape=[4, 6, 2], maxval=20, dtype=tf.int32)
print(inp)

array([[[14,  8],
    [18, 10],
    [ 6, 14],
    [ 8,  9],
    [11, 11],
    [14, 13]],

   [[ 7, 18],
    [ 4, 10],
    [15,  6],
    [ 6,  2],
    [19, 11],
    [10,  4]],

   [[ 8,  1],
    [ 1,  3],
    [ 4, 17],
    [15,  7],
    [ 0,  0],
    [ 1,  4]],

   [[ 5,  0],
    [15, 12],
    [ 1, 16],
    [ 3, 17],
    [14, 17],
    [ 2, 18]]], dtype=int32)>

だから私が正しく理解していれば、それぞれについてinp[i, :, :]：

    [[14,  8],
    [18, 10],
    [ 6, 14],
    [ 8,  9],
    [11, 11],
    [14, 13]]

最大数を含むアイテム（この場合は2行目）を保持する必要があります[18, 10]。私がすることは、最初に最後の軸に沿って最大数を見つけることです：

am = tf.math.reduce_max(inp, axis=2)
am[0, :, :]
[14,
 18,
 14,
 9,
 11,
 14]

次に、最大数を含む行のインデックスを見つけます。

am = tf.math.argmax(am, axis=1)

これらはj必要なものになり、tf.gather_ndそれらの値を取得するために使用および列挙できます。

# [*enumerate(am)] = [(0, am[0]), (1, am[1]), ...]
tf.gather_nd(inp, [*enumerate(am)])

<tf.Tensor: shape=(4, 2), dtype=int32, numpy=
array([[18, 10],
       [19, 11],
       [ 4, 17],
       [ 2, 18]], dtype=int32)>