python之在 TensorFlow 中，如何有选择地将函数仅应用于张量每一行的最大元素

我有一个二维张量，我只想对每行的最大元素应用一个函数。例如，如果我的张量是:

[[0.5, 0.7], 
 [0.4, 0.3]] 

我的函数是乘以2，那么我希望我的结果是

[[0.5, 1.4], 
 [0.8, 0.3]]. 

我目前的做法是使用tf.select，像这样:

tf.select(tf.equal(my_tensor, 
                   tf.reduce_max(my_tensor,axis=1)),  
          my_function(my_tensor), 
          my_tensor) 

想法是使用 tf.reduce_max 创建一个张量，它是每行的最大值，然后使用 tf.equal 创建一个 bool 张量，表示输入的每个元素是否等于其行最大值。然后将其作为掩码传递给 tf.select 以确定是否应应用该函数。然而，这似乎不起作用——最大张量与输入张量的维度不同，我认为广播存在一些问题。此外，如果存在两个相等的最大元素，此方法可能会遇到问题。我宁愿有一个自动打破关系的方法，也许使用 tf.arg_max。

请您参考如下方法:

我能够通过对 tf.reduce_max 的结果使用 tf.expand_dims 来修改问题中的代码，以获得要广播的正确形状的张量以匹配输入张量。结果如下所示:

tf.select(tf.equal(my_tensor, 
                   tf.expand_dims(tf.reduce_max(my_tensor,axis=1), axis=1)), 
          my_function(my_tensor), 
          my_tensor) 

虽然这可行，但我怀疑有更简洁的方法。