title: how to use for loop in tensorflow/keras when using tf.keras functional API?
body: <p>I wish to modify the CuDNNGRU Output with for loop. However, it seems like I can't do so due to the tf.GradientTape graph mode. How can I modify the CuDNNGRU in Functional API? I know that normally we can perform some matrix ops on the functional API with tf.keras.backend.* functions like K.backend.batch_dot etc. However, I have to do some complex ops like triple for loop or more, etc., If someone knows how to do so, please help! </p>
  
  <pre><code>.....source code
  x = L.Lambda(lambda fm: tf.squeeze(fm, axis=1))(x)
  gru_1 = CuDNNGRU(512, return_sequences=True, name='gru1')(x)
  gru_1b = CuDNNGRU(512, return_sequences=True, go_backwards=True,name='gru1_b')(x)
  for i in gru_1:
      .....apply some function to gru_1 outputs
  </code></pre>
  
  <p>By the way, I currently try to modify the GRU outputs with below code.</p>
  
  <pre><code>def attention(inputs):
      transpose_input = tf.transpose(inputs,perm=[0,2,1])
      atten_w = K.backend.batch_dot(inputs,transpose_input)
      atten_w = tf.linalg.set_diag(atten_w,tf.zeros(tf.shape(atten_w)[0:-1],dtype=tf.float32))
      atten_w = tf.nn.softmax(atten_w,axis=1)
      atten_v = tf.py_function(calculate_atten,inp=[inputs,atten_w],Tout=[tf.float32])
      atten_v = tf.convert_to_tensor(atten_v)
      atten_v.set_shape(self.input_shapex)
  
  
  def calculate_atten(data,atten_w):
      input_vector = data.numpy()
      atten_vectors = atten_w.numpy()
      all_batch = []
      for index,one_batch in enumerate(input_vector):
          tmp_w = atten_vectors[index]
          all_vector = []
          for j,vector in enumerate(one_batch):
              tmp = np.zeros(input_vector.shape[2])
              for w in tmp_w[j]:
                  tmp += vector*w
              all_vector.append(tmp)
          all_batch.append(all_vector)
      return all_batch
  </code></pre>
  
  <p>However, the above code, the tf.py_function return [time,features] instead of [batch,time,features], if this can be done, i can use tf.py_function to build a layer. But it seems like can't, HELP!!!!</p>
  
  <h1>Update</h1>
  
  <p>I have been able to achieve the ops, with nested tf.map_fn. Although it needs to think properly what to pass to tf.map_fn(multiple inputs must be return in multiple output). Hope this can help others</p>
  
  <pre><code>def attn_transformation(inputs):
      inputs_transpose = tf.transpose(inputs)
      atten_w = tf.tensordot(inputs,inputs_transpose,axes=1)
      def transform(data):
          multiply_data = data[0]*data[1][...,tf.newaxis]
          return [multiply_data,data[1]]
      data = tf.map_fn(lambda x:transform(x),elems=([inputs,atten_w]))
      data = tf.reduce_sum(data[0],axis=1)
      return data
  
  gru_1 = CuDNNGRU(512, return_sequences=True, name='gru1')(x)
  gru_1b = CuDNNGRU(512, return_sequences=True, go_backwards=True,name='gru1_b')(x)
  atten_vf = L.Lambda(lambda x: tf.map_fn(attn_transformation,x))(gru_1)
  </code></pre>