def test_bad_parent_values_on_gpu():
# (max_time = 4, batch_size = 1, beams = 3)
# bad parent in beam 1 time 1; appears as a negative index at time 0
end_token = 10
step_ids = _transpose_batch_time([[[1, 2, 3], [4, 5, 6], [7, 8, 9],
[-1, -1, -1]]])
parent_ids = _transpose_batch_time([[[0, 0, 0], [0, -1, 1], [2, 1, 2],
[-1, -1, -1]]])
max_sequence_lengths = [3]
expected_result = _transpose_batch_time([[[2, -1, 2], [6, 5, 6], [7, 8, 9],
[10, 10, 10]]])
if options.TF_ADDONS_PY_OPS:
# The Python version has the same behavior on CPU and GPU.
with pytest.raises(tf.errors.InvalidArgumentError, match="parent id"):
_ = gather_tree(
step_ids=step_ids,
parent_ids=parent_ids,
max_sequence_lengths=max_sequence_lengths,
end_token=end_token,
)
else:
beams = gather_tree(
step_ids=step_ids,
parent_ids=parent_ids,
max_sequence_lengths=max_sequence_lengths,
end_token=end_token,
)
np.testing.assert_equal(expected_result, beams.numpy())
def test_gather_tree():
# (max_time = 3, batch_size = 2, beam_width = 3)
# create (batch_size, max_time, beam_width) matrix and transpose it
predicted_ids = np.array(
[[[1, 2, 3], [4, 5, 6], [7, 8, 9]], [[2, 3, 4], [5, 6, 7], [8, 9, 10]]
],
dtype=np.int32,
).transpose([1, 0, 2])
parent_ids = np.array(
[[[0, 0, 0], [0, 1, 1], [2, 1, 2]], [[0, 0, 0], [1, 2, 0], [2, 1, 1]]],
dtype=np.int32,
).transpose([1, 0, 2])
# sequence_lengths is shaped (batch_size = 3)
max_sequence_lengths = [3, 3]
expected_result = np.array([[[2, 2, 2], [6, 5, 6], [7, 8, 9]],
[[2, 4, 4], [7, 6, 6],
[8, 9, 10]]]).transpose([1, 0, 2])
res = gather_tree(
predicted_ids,
parent_ids,
max_sequence_lengths=max_sequence_lengths,
end_token=11,
)
np.testing.assert_equal(expected_result, res)
def testBadParentValuesOnGPU(self):
# Only want to run this test on CUDA devices, as gather_tree is not
# registered for SYCL devices.
if not tf.test.is_gpu_available(cuda_only=True):
return
# (max_time = 4, batch_size = 1, beams = 3)
# bad parent in beam 1 time 1; appears as a negative index at time 0
end_token = 10
step_ids = _transpose_batch_time(
[[[1, 2, 3], [4, 5, 6], [7, 8, 9], [-1, -1, -1]]]
)
parent_ids = _transpose_batch_time(
[[[0, 0, 0], [0, -1, 1], [2, 1, 2], [-1, -1, -1]]]
)
max_sequence_lengths = [3]
expected_result = _transpose_batch_time(
[[[2, -1, 2], [6, 5, 6], [7, 8, 9], [10, 10, 10]]]
)
with tf.device("/device:GPU:0"):
beams = gather_tree(
step_ids=step_ids,
parent_ids=parent_ids,
max_sequence_lengths=max_sequence_lengths,
end_token=end_token,
)
self.assertAllEqual(expected_result, self.evaluate(beams))
def testGatherTreeOne(self):
# (max_time = 4, batch_size = 1, beams = 3)
end_token = 10
step_ids = _transpose_batch_time([[[1, 2, 3], [4, 5, 6], [7, 8, 9],
[-1, -1, -1]]])
parent_ids = _transpose_batch_time([[[0, 0, 0], [0, 1, 1], [2, 1, 2],
[-1, -1, -1]]])
max_sequence_lengths = [3]
expected_result = _transpose_batch_time([[[2, 2, 2], [6, 5, 6],
[7, 8, 9], [10, 10, 10]]])
beams = gather_tree(step_ids=step_ids,
parent_ids=parent_ids,
max_sequence_lengths=max_sequence_lengths,
end_token=end_token)
with self.cached_session(use_gpu=True):
self.assertAllEqual(expected_result, self.evaluate(beams))
def testGatherTreeBatch(self):
batch_size = 10
beam_width = 15
max_time = 8
max_sequence_lengths = [0, 1, 2, 4, 7, 8, 9, 10, 11, 0]
end_token = 5
with self.cached_session(use_gpu=True):
step_ids = np.random.randint(0,
high=end_token + 1,
size=(max_time, batch_size,
beam_width))
parent_ids = np.random.randint(0,
high=beam_width - 1,
size=(max_time, batch_size,
beam_width))
beams = gather_tree(
step_ids=step_ids.astype(np.int32),
parent_ids=parent_ids.astype(np.int32),
max_sequence_lengths=max_sequence_lengths,
end_token=end_token,
)
self.assertEqual((max_time, batch_size, beam_width), beams.shape)
beams_value = self.evaluate(beams)
for b in range(batch_size):
# Past max_sequence_lengths[b], we emit all end tokens.
b_value = beams_value[max_sequence_lengths[b]:, b, :]
self.assertAllClose(b_value, end_token * np.ones_like(b_value))
for batch, beam in itertools.product(range(batch_size),
range(beam_width)):
v = np.squeeze(beams_value[:, batch, beam])
if end_token in v:
found_bad = np.where(v == -1)[0]
self.assertEqual(0, len(found_bad))
found = np.where(v == end_token)[0]
found = found[0] # First occurrence of end_token.
# If an end_token is found, everything before it should be a
# valid id and everything after it should be -1.
if found > 0:
self.assertAllEqual(
v[:found - 1] >= 0,
np.ones_like(v[:found - 1], dtype=bool),
)
self.assertAllClose(
v[found + 1:], end_token * np.ones_like(v[found + 1:]))
def test_bad_parent_values_on_cpu():
# (batch_size = 1, max_time = 4, beams = 3)
# bad parent in beam 1 time 1
end_token = 10
step_ids = _transpose_batch_time([[[1, 2, 3], [4, 5, 6], [7, 8, 9],
[-1, -1, -1]]])
parent_ids = _transpose_batch_time([[[0, 0, 0], [0, -1, 1], [2, 1, 2],
[-1, -1, -1]]])
max_sequence_lengths = [3]
with pytest.raises(tf.errors.InvalidArgumentError, match="parent id"):
_ = gather_tree(
step_ids=step_ids,
parent_ids=parent_ids,
max_sequence_lengths=max_sequence_lengths,
end_token=end_token,
)
def test_gather_tree_one():
# (max_time = 4, batch_size = 1, beams = 3)
end_token = 10
step_ids = _transpose_batch_time([[[1, 2, 3], [4, 5, 6], [7, 8, 9],
[-1, -1, -1]]])
parent_ids = _transpose_batch_time([[[0, 0, 0], [0, 1, 1], [2, 1, 2],
[-1, -1, -1]]])
max_sequence_lengths = [3]
expected_result = _transpose_batch_time([[[2, 2, 2], [6, 5, 6], [7, 8, 9],
[10, 10, 10]]])
beams = gather_tree(
step_ids=step_ids,
parent_ids=parent_ids,
max_sequence_lengths=max_sequence_lengths,
end_token=end_token,
)
np.testing.assert_equal(expected_result, beams.numpy())
def testBadParentValuesOnCPU(self):
# (batch_size = 1, max_time = 4, beams = 3)
# bad parent in beam 1 time 1
end_token = 10
step_ids = _transpose_batch_time([[[1, 2, 3], [4, 5, 6], [7, 8, 9],
[-1, -1, -1]]])
parent_ids = _transpose_batch_time([[[0, 0, 0], [0, -1, 1], [2, 1, 2],
[-1, -1, -1]]])
max_sequence_lengths = [3]
with tf.device("/cpu:0"):
msg = r"parent id -1 at \(batch, time, beam\) == \(0, 0, 1\)"
with self.assertRaisesOpError(msg):
beams = gather_tree(step_ids=step_ids,
parent_ids=parent_ids,
max_sequence_lengths=max_sequence_lengths,
end_token=end_token)
self.evaluate(beams)
def test_bad_parent_values_on_gpu():
# (max_time = 4, batch_size = 1, beams = 3)
# bad parent in beam 1 time 1; appears as a negative index at time 0
end_token = 10
step_ids = _transpose_batch_time([[[1, 2, 3], [4, 5, 6], [7, 8, 9],
[-1, -1, -1]]])
parent_ids = _transpose_batch_time([[[0, 0, 0], [0, -1, 1], [2, 1, 2],
[-1, -1, -1]]])
max_sequence_lengths = [3]
expected_result = _transpose_batch_time([[[2, -1, 2], [6, 5, 6], [7, 8, 9],
[10, 10, 10]]])
beams = gather_tree(
step_ids=step_ids,
parent_ids=parent_ids,
max_sequence_lengths=max_sequence_lengths,
end_token=end_token,
)
np.testing.assert_equal(expected_result, beams.numpy())
def test_gather_tree_batch():
batch_size = 10
beam_width = 15
max_time = 8
max_sequence_lengths = [0, 1, 2, 4, 7, 8, 9, 10, 11, 0]
end_token = 5
step_ids = np.random.randint(0,
high=end_token + 1,
size=(max_time, batch_size, beam_width))
parent_ids = np.random.randint(0,
high=beam_width - 1,
size=(max_time, batch_size, beam_width))
beams = gather_tree(
step_ids=step_ids.astype(np.int32),
parent_ids=parent_ids.astype(np.int32),
max_sequence_lengths=max_sequence_lengths,
end_token=end_token,
)
beams = beams.numpy()
assert (max_time, batch_size, beam_width) == beams.shape
for b in range(batch_size):
# Past max_sequence_lengths[b], we emit all end tokens.
b_value = beams[max_sequence_lengths[b]:, b, :]
np.testing.assert_allclose(b_value, end_token * np.ones_like(b_value))
for batch, beam in itertools.product(range(batch_size), range(beam_width)):
v = np.squeeze(beams[:, batch, beam])
if end_token in v:
found_bad = np.where(v == -1)[0]
assert 0 == len(found_bad)
found = np.where(v == end_token)[0]
found = found[0] # First occurrence of end_token.
# If an end_token is found, everything before it should be a
# valid id and everything after it should be end_token.
if found > 0:
np.testing.assert_equal(
v[:found - 1] >= 0,
np.ones_like(v[:found - 1], dtype=bool),
)
np.testing.assert_allclose(v[found + 1:],
end_token * np.ones_like(v[found + 1:]))
