docstr: tf.function compilation correctly handles if statement.
docstr_tok: ["tf.function","compilation","correctly","handles","if","statement","."]
code: def test_with_tf_function_with_booleans(self, autograph):
      """tf.function compilation correctly handles if statement."""
  
      layer = dropout.Dropout(rate=0.5)
      layer = tf.function(layer, autograph=autograph)
  
      inputs = tf.ones([2, 5, 3, 3, 3])
      expected = tf.zeros_like(inputs)
  
      for is_training in (True, False):
          outputs = layer(inputs, is_training)
          self.assertEqual(outputs.shape, expected.shape)
  
file_path: sonnet/sonnet/src/dropout_test.py
imports: from absl.testing import parameterized
  import numpy
  from sonnet.src import dropout
  from sonnet.src import test_utils
  import tensorflow
label_desc: Use image node_id 4 for calling the DropoutTest obj's underlying member method code with example usage: obj.test_with_tf_function_with_booleans(autograph) without return types