def test_sequence_mask(self):
with self.test_session(graph=tf.Graph()):
lengths = tf.constant([1, 2, 3])
mask = tensor_utils.sequence_mask(lengths)
transposed_mask = tensor_utils.sequence_mask(lengths,
transpose=True)
self.assertAllEqual(mask.eval(),
[[True, False, False], [True, True, False],
[True, True, True]]) # pyformat: disable
self.assertAllEqual(transposed_mask.eval(),
[[True, True, True], [False, True, True],
[False, False, True]]) # pyformat: disable
def test_sequence_mask_with_dtype(self):
with self.test_session(graph=tf.Graph()):
lengths = tf.constant([1, 2, 3])
mask = tensor_utils.sequence_mask(lengths, dtype=tf.int32)
transposed_mask = tensor_utils.sequence_mask(lengths,
dtype=tf.int32,
transpose=True)
self.assertAllEqual(
mask.eval(),
[[1, 0, 0], [1, 1, 0], [1, 1, 1]]) # pyformat: disable
self.assertAllEqual(
transposed_mask.eval(),
[[1, 1, 1], [0, 1, 1], [0, 0, 1]]) # pyformat: disable
