code: def test_static_scalar_condition(self):
      fn_calls = [0, 0]
      ones = tf.ones([10])
      zeros = tf.zeros([10])
  
      def fn1():
          fn_calls[0] += 1
          return ones
  
      def fn2():
          fn_calls[1] += 1
          return zeros
  
      self.assertAllEqual(zeros, ps.smart_where(False, fn1, fn2))
      self.assertEqual([0, 1], fn_calls)
      self.assertAllEqual(ones, ps.smart_where(True, fn1, fn2))
      self.assertEqual([1, 1], fn_calls)
      self.assertAllEqual(
          zeros, ps.smart_where(tf.constant(False), fn1, fn2)
      )
      self.assertEqual([1, 2], fn_calls)
      self.assertAllEqual(
          ones, ps.smart_where(tf.constant(True), fn1, fn2)
      )
      self.assertEqual([2, 2], fn_calls)
      self.assertAllEqual(
          zeros, ps.smart_where(np.array(False), fn1, fn2)
      )
      self.assertEqual([2, 3], fn_calls)
      self.assertAllEqual(
          ones, ps.smart_where(np.array(True), fn1, fn2)
      )
      self.assertEqual([3, 3], fn_calls)
  
      self.assertAllEqual(zeros, ps.smart_where(0, fn1, fn2))
      self.assertEqual([3, 4], fn_calls)
      self.assertAllEqual(ones, ps.smart_where(1, fn1, fn2))
      self.assertEqual([4, 4], fn_calls)
      self.assertAllEqual(
          zeros, ps.smart_where(tf.constant(0), fn1, fn2)
      )
      self.assertEqual([4, 5], fn_calls)
      self.assertAllEqual(
          ones, ps.smart_where(tf.constant(1), fn1, fn2)
      )
      self.assertEqual([5, 5], fn_calls)
      self.assertAllEqual(zeros, ps.smart_where(np.array(0), fn1, fn2))
      self.assertEqual([5, 6], fn_calls)
      self.assertAllEqual(ones, ps.smart_where(np.array(1), fn1, fn2))
      self.assertEqual([6, 6], fn_calls)
  
imports: from absl.testing import parameterized
  import numpy
  import tensorflow.compat.v1
  import tensorflow.compat.v2
  from tensorflow_probability.python.internal import dtype_util
  from tensorflow_probability.python.internal import prefer_static
  from tensorflow_probability.python.internal import test_util
label_desc: Use image node_id 1 for calling the SmartWhereTest obj's underlying member method code with example usage: obj.test_static_scalar_condition() and returns: ones, zeros