def testSegmentedRange(self):
# TODO(b/157666350): Turn `testSegmentedRange` into a hypothesis test.
limits = np.array([3, 1, 4, 5, 4, 3, 10, 1])
self.assertAllEqual(self.evaluate(tf.ragged.range(limits).flat_values),
bates._segmented_range(limits))
def testSegmentedRange(self, limits):
limits = np.int32(limits)
self.assertAllEqual(
tf.ragged.range(limits).flat_values,
bates._segmented_range(limits))