def test_pack_sequences_k_bfloat16(self):
inputs = tf.convert_to_tensor([
[1, 2, 3],
[4, 5, 0],
[6, 0, 0],
],
dtype=tf.bfloat16)
max_length = tf.convert_to_tensor(5, dtype=tf.int32)
(packed, segmentation,
position) = pack_sequences_ops.pack_sequences_k([inputs],
[max_length])
(inputs_packed, ) = packed
(inputs_segmentation, ) = segmentation
(inputs_position, ) = position
self.assertAllEqual(inputs_packed, [
[1, 2, 3, 4, 5],
[6, 0, 0, 0, 0],
])
self.assertEqual(inputs_packed.dtype, tf.bfloat16)
self.assertAllEqual(inputs_segmentation, [
[1, 1, 1, 2, 2],
[1, 0, 0, 0, 0],
])
self.assertAllEqual(inputs_position, [
[0, 1, 2, 0, 1],
[0, 0, 0, 0, 0],
])
def _pack_sequences_k(inputs, targets, input_max_length, target_max_length):
"""Wrapper for pack_sequences_k with same interface as pack_sequences_2."""
inputs = tf.convert_to_tensor(inputs, tf.int32)
targets = tf.convert_to_tensor(targets, tf.int32)
input_max_length = tf.convert_to_tensor(input_max_length, dtype=tf.int32)
target_max_length = tf.convert_to_tensor(target_max_length, dtype=tf.int32)
(packed, segmentation, position) = pack_sequences_ops.pack_sequences_k(
[inputs, targets], [input_max_length, target_max_length])
(inputs_packed, targets_packed) = packed
(inputs_segmentation, targets_segmentation) = segmentation
(inputs_position, targets_position) = position
return (inputs_packed, inputs_segmentation, inputs_position,
targets_packed, targets_segmentation, targets_position)
def test_pack_sequences_k(self):
inputs = tf.convert_to_tensor([
[1, 2, 3],
[4, 5, 0],
[6, 0, 0],
],
dtype=tf.int32)
targets = tf.convert_to_tensor([
[10, 0, 0],
[20, 30, 40],
[50, 60, 0],
],
dtype=tf.int32)
max_length = tf.convert_to_tensor(5, dtype=tf.int32)
(packed, segmentation, position) = pack_sequences_ops.pack_sequences_k(
[inputs, targets], [max_length, max_length])
(inputs_packed, targets_packed) = packed
(inputs_segmentation, targets_segmentation) = segmentation
(inputs_position, targets_position) = position
self.assertAllEqual(inputs_packed, [
[1, 2, 3, 4, 5],
[6, 0, 0, 0, 0],
])
self.assertAllEqual(inputs_segmentation, [
[1, 1, 1, 2, 2],
[1, 0, 0, 0, 0],
])
self.assertAllEqual(inputs_position, [
[0, 1, 2, 0, 1],
[0, 0, 0, 0, 0],
])
self.assertAllEqual(targets_packed, [
[10, 20, 30, 40, 0],
[50, 60, 0, 0, 0],
])
self.assertAllEqual(targets_segmentation, [
[1, 2, 2, 2, 0],
[1, 1, 0, 0, 0],
])
self.assertAllEqual(targets_position, [
[0, 0, 1, 2, 0],
[0, 1, 0, 0, 0],
])
def test_pack_sequences_k_multi_input(self):
input_tokens = tf.convert_to_tensor([
[1, 2, 3],
[4, 5, 0],
[6, 0, 0],
],
dtype=tf.int32)
input_vectors = tf.convert_to_tensor([
[[0, 1, 2], [1, 2, 3], [2, 3, 4]],
[[3, 4, 5], [4, 5, 6], [0, 0, 0]],
[[5, 6, 7], [0, 0, 0], [0, 0, 0]],
],
dtype=tf.float32)
targets = tf.convert_to_tensor([
[10, 0, 0],
[20, 30, 40],
[50, 60, 0],
],
dtype=tf.int32)
(packed, segmentation, position) = pack_sequences_ops.pack_sequences_k(
[input_tokens, input_vectors, targets], [5, 3, 5])
(input_tokens_packed, input_vectors_packed, targets_packed) = packed
(input_tokens_segmentation, input_vectors_segmentation,
targets_segmentation) = segmentation
(input_tokens_position, input_vectors_position,
targets_position) = position
self.assertAllEqual(input_tokens_packed, [
[1, 2, 3, 0, 0],
[4, 5, 6, 0, 0],
])
self.assertAllEqual(input_vectors_packed, [
[[0, 1, 2], [1, 2, 3], [2, 3, 4]],
[[3, 4, 5], [4, 5, 6], [5, 6, 7]],
])
self.assertAllEqual(input_tokens_segmentation, [
[1, 1, 1, 0, 0],
[1, 1, 2, 0, 0],
])
self.assertAllEqual(input_vectors_segmentation, [
[1, 1, 1],
[1, 1, 2],
])
self.assertAllEqual(input_tokens_position, [
[0, 1, 2, 0, 0],
[0, 1, 0, 0, 0],
])
self.assertAllEqual(input_vectors_position, [
[0, 1, 2],
[0, 1, 0],
])
self.assertAllEqual(targets_packed, [
[10, 0, 0, 0, 0],
[20, 30, 40, 50, 60],
])
self.assertAllEqual(targets_segmentation, [
[1, 0, 0, 0, 0],
[1, 1, 1, 2, 2],
])
self.assertAllEqual(targets_position, [
[0, 0, 0, 0, 0],
[0, 1, 2, 0, 1],
])