code: def test_layout_prop_v2_if(self):
      a = constant_op.constant([0, 1, 2, 1], dtype=dtypes.float32)
      a = numpy_util.pack_numpy(a, self.unsharded_layout)
  
      @polymorphic_function.function
      def function_with_if(t):
          if math_ops.equal(math_ops.reduce_sum(t), 0):
              t = math_ops.sqrt(t)
              return api.relayout(t, self.x_layout)
          else:
              return array_ops.zeros_like_v2(t)
  
      dtensor_result = function_with_if(a)
      api.check_layout(dtensor_result, self.x_layout)
  
imports: from absl.testing import parameterized
  import numpy
  from tensorflow.dtensor.python import api
  from tensorflow.dtensor.python import d_variable
  from tensorflow.dtensor.python import layout
  from tensorflow.dtensor.python import numpy_util
  from tensorflow.dtensor.python.tests import test_util
  from tensorflow.python.eager import backprop
  from tensorflow.python.eager.polymorphic_function import polymorphic_function
  from tensorflow.python.framework import constant_op
  from tensorflow.python.framework import dtypes
  from tensorflow.python.ops import array_ops
  from tensorflow.python.ops import gen_nn_ops
  from tensorflow.python.ops import math_ops
  from tensorflow.python.ops import nn_ops
  from tensorflow.python.ops import stateless_random_ops
  from tensorflow.python.ops import variables
  from tensorflow.python.platform import test
label_desc: Use image node_id 7 for calling the LayoutPropagationV2Test obj's underlying member method code with example usage: obj.test_layout_prop_v2_if() and returns: api, array_ops