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

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

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