def testInputShapeChecks(self):
    """Input shape checks can catch some, but not all, shape problems."""
    # 1. Inputs have incorrect or incompatible ranks:
    for args in [dict(values=[[5, 5]], indices=1),
                 dict(values=[5, 5], indices=[1]),
                 dict(values=[[[5, 5]]], indices=[1]),
                 dict(values=[[5, 5]], indices=[[[1]]]),]:
      with self.assertRaisesRegexp(ValueError, "do not correspond"):
        indexing_ops.batched_index(**args)

    # 2. Inputs have correct, compatible ranks but incompatible sizes:
    for args in [dict(values=[[5, 5]], indices=[1, 1]),
                 dict(values=[[5, 5], [5, 5]], indices=[1]),
                 dict(values=[[[5, 5], [5, 5]]], indices=[[1, 1], [1, 1]]),
                 dict(values=[[[5, 5], [5, 5]]], indices=[[1], [1]]),]:
      with self.assertRaisesRegexp(ValueError, "incompatible shapes"):
        indexing_ops.batched_index(**args)

    # (Correct ranks and sizes work fine, though):
    indexing_ops.batched_index(
        values=[[5, 5]], indices=[1])
    indexing_ops.batched_index(
        values=[[[5, 5], [5, 5]]], indices=[[1, 1]])

    # 3. Shape-checking works with fully-specified placeholders, or even
    # partially-specified placeholders that still provide evidence of having
    # incompatible shapes or incorrect ranks.
    for sizes in [dict(q_size=[4, 3], a_size=[4, 1]),
                  dict(q_size=[4, 2, 3], a_size=[4, 1]),
                  dict(q_size=[4, 3], a_size=[5, None]),
                  dict(q_size=[None, 2, 3], a_size=[4, 1]),
                  dict(q_size=[4, 2, 3], a_size=[None, 1]),
                  dict(q_size=[4, 2, 3], a_size=[5, None]),
                  dict(q_size=[None, None], a_size=[None, None]),
                  dict(q_size=[None, None, None], a_size=[None]),]:
      with self.assertRaises(ValueError):
        indexing_ops.batched_index(
            tf.placeholder(tf.float32, sizes["q_size"]),
            tf.placeholder(tf.int32, sizes["a_size"]))

    # But it can't work with 100% certainty if full shape information is not
    # known ahead of time. These cases generate no errors; some make warnings:
    for sizes in [dict(q_size=None, a_size=None),
                  dict(q_size=None, a_size=[4]),
                  dict(q_size=[4, 2], a_size=None),
                  dict(q_size=[None, 2], a_size=[None]),
                  dict(q_size=[None, 2, None], a_size=[None, 2]),
                  dict(q_size=[4, None, None], a_size=[4, None]),
                  dict(q_size=[None, None], a_size=[None]),
                  dict(q_size=[None, None, None], a_size=[None, None]),]:
      indexing_ops.batched_index(
          tf.placeholder(tf.float32, sizes["q_size"]),
          tf.placeholder(tf.int32, sizes["a_size"]))

    # And it can't detect invalid indices at construction time, either.
    indexing_ops.batched_index(values=[[5, 5, 5]], indices=[1000000000])
Input shape checks can catch some, but not all, shape problems.
trfl/indexing_ops_test.py