def test_where_select_elements(self):
    """Test for where() when selecting individual elements."""
    with self.test_session(graph=tf.Graph()):
      condition = tf.constant([[True, False],
                               [False, True]])  # pyformat: disable
      if_true = tf.constant([[1, 2],
                             [3, 4]])  # pyformat: disable
      if_false = tf.constant([[5, 6],
                              [7, 8]])  # pyformat: disable

      result = tensor_utils.where(condition, if_true, if_false)

      self.assertAllEqual(result.eval(),
                          [[1, 6],
                           [7, 4]])  # pyformat: disable
Test for where() when selecting individual elements.
language/common/utils/tensor_utils_test.py