def _ensure_no_eos(feat, v):
if feat not in self.output_features:
return v
with tf.control_dependencies([
tf.assert_none_equal(
v, tf.constant(1, tf.int64),
message="Feature '{feat}' unexpectedly contains EOS=1 token "
"after {label}.".format(feat=feat, label=error_label))
]):
return v
def _finalize(self, _, contents):
"""Structure output and compute segment and position metadata."""
# The output shape information is lost during the filter; however we can
# guarantee the shape. (That's the point of this exercise, after all!)
contents.set_shape((self._packed_length, self._num_sequences * 2))
# Both the dummy branch of the scan step function and the eviction dataset
# use vectors of minus one. The cost of this check is negligible and the
# leakage of such dummy sequences would be difficult to debug downstream.
check_leaks = tf.assert_none_equal(contents, -tf.ones_like(contents))
with tf.control_dependencies([check_leaks]):
contents = tf.identity(contents)
segment, position = self._compute_auxiliary_structure(contents)
return {"contents": contents[:, :self._num_sequences],
"segment": segment, "position": position}
