def _repeat_molecules(self, molecules: tf.Tensor,
                        char_seq_length: tf.Tensor) -> tf.Tensor:
    """Repeats molecules to make them the same length as the char sequence."""

    rate = self.config.downsampling_rate

    molecules_without_extra_cls = molecules[:, 1:, :]
    # `repeated`: [batch_size, almost_char_seq_len, molecule_hidden_size]
    repeated = tf.repeat(molecules_without_extra_cls, repeats=rate, axis=-2)

    # So far, we've repeated the elements sufficient for any `char_seq_length`
    # that's a multiple of `downsampling_rate`. Now we account for the last
    # n elements (n < `downsampling_rate`), i.e. the remainder of floor
    # division. We do this by repeating the last molecule a few extra times.
    last_molecule = molecules[:, -1:, :]
    remainder_length = tf.floormod(char_seq_length, rate)
    remainder_repeated = tf.repeat(
        last_molecule,
        # +1 molecule to compensate for truncation.
        repeats=remainder_length + rate,
        axis=-2)

    # `repeated`: [batch_size, char_seq_len, molecule_hidden_size]
    return tf.concat([repeated, remainder_repeated], axis=-2)
Repeats molecules to make them the same length as the char sequence.
language/canine/modeling.py