[Question] Pool bert subwords back to word level?

We currently have some code that does the bert op in the graph. Do you have a method of pooling the bert tokens back to the word level? Just curious if you guys had an idiomatic method of doing this.

Currently the word level is flattened out in the call to merge_dims()

However it would be nice to merge the subword vectors (after sending through bert) back to word level, via some pooling operation.

         tokens = self.tokenizer.tokenize(raw_text)

        # the tokenizer produces subword level ragged tensors
        # these need to be merged back to be word level per utterance
        # merge_dims() flattens [[[]]] -> [[]]

        # Trim the ragged tokens to max_seq_len - 2 (to account for CLS/SEP)
        ragged_tokens = tokens.merge_dims(
            inner_axis=2, outer_axis=1)[:, :self.max_seq_len - 2]
        ragged_tokens = tf.cast(ragged_tokens, tf.int32)

        # Concat CLS/SEP before conversion to sparse
        cls_tokens = tf.reshape(
            tf.tile([_CLS_ID], [tokens.nrows()]), [tokens.nrows(), 1])
        sep_tokens = tf.reshape(
            tf.tile([_SEP_ID], [tokens.nrows()]), [tokens.nrows(), 1])
        # add CLS and SEP to start and end
        ragged_tokens = tf.concat([cls_tokens,
                                   ragged_tokens,
                                   sep_tokens], axis=1)

        # to dense, fill in 0 with _PAD
        input_word_ids = ragged_tokens.to_tensor(default_value=_PAD_ID)

        paddings = [[0, 0],
                    [0, self.max_seq_len - tf.shape(input_word_ids)[1]]]

        input_word_ids = tf.pad(input_word_ids, paddings,
                                'CONSTANT', constant_values=_PAD_ID)

        # calculate the input masks and cast
        input_mask = tf.where((input_word_ids == _PAD_ID) |
                              (input_word_ids == _CLS_ID) |
                              (input_word_ids == _SEP_ID),
                              0,
                              tf.ones(self.max_seq_len, tf.int32))

        # calculate the segment ids
        segment_ids = tf.cast(input_word_ids > 0, tf.int32)