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}