def METHOD_NAME(self):
    positive_matches = tf.constant(
        [
            [
                True,
                False,
                False,
                False,
                False,
                False,
                False,
                False,
                False,
                False,
            ],
            [
                False,
                False,
                False,
                False,
                False,
                False,
                True,
                False,
                False,
                False,
            ],
        ]
    )
    negative_matches = tf.constant(
        [
            [False, True, True, True, True, True, True, True, True, True],
            [True, True, True, True, True, True, False, True, True, True],
        ]
    )
    num_samples = 5
    positive_fraction = 0.2
    res = balanced_sample(
        positive_matches, negative_matches, num_samples, positive_fraction
    )
    # the 1st element from the 1st batch must be selected, given it's the
    # only one
    self.assertAllClose(res[0][0], 1)
    # the 7th element from the 2nd batch must be selected, given it's the
    # only one
    self.assertAllClose(res[1][6], 1)