How to convert sample program in tensorflow efficiently?-CodePudding

My Code :

  data = [0, 2]
  f = numpy.array([[[1, 2], [3, 4]],
       [[4, 5], [7, 5]], 
       [[6, 3], [7, 9]]])
  l = []
  for i in data :
    l.append(f[i])
  return np.maximum.reduce(l)

Output : [[6, 3], [7, 9]] Element wise maximum between f[0] and f[2] as data is 0 and 2
All i need is to implement the same code in tensorflow format using tf.while_loop and any other tensorflow function

CodePudding user response：

You can try something like this:

import tensorflow as tf

data = [0, 2]
f = tf.constant([[[1, 2], [3, 4]],
      [[4, 5], [7, 5]], 
      [[6, 3], [7, 9]]])
x = tf.gather(f, data)
x = tf.reduce_max(x, axis=0)
print(x)

tf.Tensor(
[[6 3]
 [7 9]], shape=(2, 2), dtype=int32)

Regarding your question in the comments, try something like this:

fn = 4 
x = tf.random.normal((1, 2, 2, 4))
x = tf.squeeze(tf.split(x[0, :, :, :], fn, axis=-1), axis=-1)