def testPaddingProvideShape(self):
self.assertEqual(
[[0, 0, 0], [0, 0, 0]], utils.pad([], pad_token=0, pad_shape=(2, 3))
)
self.assertEqual(
[[1, 2, 3], [1, 0, 0]],
utils.pad([[1, 2, 3], [1]], pad_token=0, pad_shape=(2, 3)),
)
self.assertEqual([], utils.pad([], pad_token=0, pad_shape=()))
def test_fp16_padding(self):
nested_lists = [[1, 2, 3], [4, 5]]
padded_lists = pad(nested_lists, pad_token=0)
expected_lists = [[1, 2, 3], [4, 5, 0]]
self.assertEqual(padded_lists, expected_lists)
precision.FP16_ENABLED = True
padded_lists = pad(nested_lists, pad_token=0)
expected_lists = [[1, 2, 3, 0, 0, 0, 0, 0], [4, 5, 0, 0, 0, 0, 0, 0]]
self.assertEqual(padded_lists, expected_lists)
precision.FP16_ENABLED = False
def testPadding(self):
self.assertEqual([[1, 2, 3], [1, 0, 0]],
utils.pad([[1, 2, 3], [1]], pad_token=0))
self.assertEqual(
[[[1], [2], [3]], [[1], [0], [0]]],
utils.pad([[[1], [2], [3]], [[1]]], pad_token=0),
)
self.assertEqual(
[[1, 2, 3, 4, 5, 6, 7], [9, 9, 9, 9, 9, 9, 9]],
utils.pad([[1, 2, 3, 4, 5, 6, 7], []], pad_token=9),
)
