I use this toy example to measure performance of if statements in the forward loop.

  
      
      colab.research.google.com
  
  
    

Google Colaboratory 9



  
  
    
    
  
  


At least for XLA devices (such as in COLAB) when the conditional statement is fixed the performance doesn’t seem to be affected (e.g. model properties). However when the conditional statement is highly dependant on input the performance is affected.
Are conditions handled in GPU or the code is optimized for GPU? Is the optimization and conversion of the code taking all of that into account?
Is there a general rule for code level optimization for IF statements or it doesn’t really matter?