code: def test_if():
      def create_if_class(_condition=True):
          class If(tf.Module):
              def get_input(self):
                  return np.ones((2, 2), dtype="float32")
  
              @tf.function(
                  input_signature=[
                      tf.TensorSpec(shape=(2, 2), dtype=tf.float32)
                  ]
              )
              def func(self, x):
                  @tf.function(
                      input_signature=[
                          tf.TensorSpec(shape=(2, 2), dtype=tf.float32)
                      ]
                  )
                  def double(x):
                      return 2 * x
  
                  @tf.function(
                      input_signature=[
                          tf.TensorSpec(shape=(2, 2), dtype=tf.float32)
                      ]
                  )
                  def triple(x):
                      return 3 * x
  
                  output = tf.raw_ops.If(
                      cond=_condition,
                      input=[x],
                      Tout=[tf.float32],
                      output_shapes=[(2, 2)],
                      then_branch=double.get_concrete_function(),
                      else_branch=triple.get_concrete_function(),
                  )
                  return output[0]
  
          return If
  
      for cond in [True, False]:
          if_class = create_if_class(_condition=cond)
          run_func_graph(if_class, runtime="vm")
          run_model_graph(if_class)
  
imports: import tempfile
  import tensorflow
  import numpy
  import pytest
  from common import compare_tf_tvm
  from common import run_tf_code
label_desc: Use image node_id 30 for calling a global function with example usage: test_if() and returns: If, np, output, unknown, unknown