title: Efficient conditional triggering of decoder in tensorflow
body: <p>I am implementing an encoder-(dual-)decoder model in tensorflow. The decoder is RNN-type. The input to the decoder is a feature map, the output of the previous time-step and the hidden state of the decoder from the previous time-step. I only only want to trigger the decoder(s) when the prediction from the previous time-step is one of a particular set of tokens.</p>
  <p>I have tried using <code>tf.boolean_mask</code> on the prediction of the previous time-step to remove those examples that do not predict a trigger-token. Below is an example:</p>
  <pre><code># initialize input
  dec_input = tf.expand_dims([token2integer['&lt;start&gt;']] * target.shape[0], 1)
  features = encoder(img_tensor)
  hidden = decoder.reset_state(batch_size=target.shape[0])
  
  # make first prediction
  predictions, hidden, _ = decoder(dec_input, features, hidden)
  
  # add to total loss
  loss += loss_function(dec_input, predictions)
  
  # construct input of next time-step (here with teacher forcing)
  dec_input =tf.expand_dims(target[:, i], -1)
  
  #compute mask to only trigger for certain predictions
  mask_struc = compute_mask_struc(dec_input) 
  
  # apply mask to input
  features = tf.boolean_mask(features, mask_struc)                    
  hidden = tf.boolean_mask(hidden, mask_struc)
  target = tf.boolean_mask(target, mask_struc ) 
  dec_input = tf.boolean_mask(dec_input, mask_struc ) 
  
  # make next prediction and so on ...
  </code></pre>
  <p>I have implemented this into a training function. My implementation is working but it is slow. And when I run the function as a graph (with @tf.function) it gets 10x slower. If I remove the <code>boolen_mask</code> and run as a graph (with @tf.function) it is faster than without the @tf.function.</p>
  <p>How can I speed up the execution (with or without the @tf.function)?</p>
  <p>My ideas:</p>
  <ul>
  <li>fix whatever is making the graph execution slow: I don't know how.</li>
  <li>find alternative approach (without <code>boolean_mask</code>): I need inspiration</li>
  <li>give up and try with PyTorch which I am more familiar with: not guaranteed to be faster.</li>
  </ul>