def combine_transition_probabilities(self, a: tf.Tensor, b: tf.Tensor) -> tf.Tensor:
        """Transforms logarithmic transition probabilities a and b.

        Args:
            a:      shape = [batch, DIMS_A, max_logit_length, max_label_length + 1, state]
            b:      shape = [batch, max_logit_length, max_label_length + 1, state, DIMS_B]

        Returns:    shape = [batch, DIMS_A, max_logit_length, num_tokens, DIMS_B]
        """
        assert len(a.shape) >= 4
        assert len(b.shape) >= 4
        assert a.shape[-1] == 2
        assert b.shape[3] == 2

        dims_a = tf.shape(a)[1:-3]
        dims_b = tf.shape(b)[4:]
        a = tf.reshape(a, shape=[self.batch_size, -1, self.max_logit_length, self.max_label_length_plus_one, 2, 1])
        # shape = [batch_size, dims_a, max_logit_length, max_label_length + 1, state, 1]
        b = tf.reshape(b, shape=[self.batch_size, 1, self.max_logit_length, self.max_label_length_plus_one, 2, -1])
        # shape = [batch_size, 1, max_logit_length, max_label_length + 1, state, dims_b]

        # Either open or closed state from alpha and only closed state from beta
        ab_term = tf.reduce_logsumexp(a, 4) + b[:, :, :, :, 0]
        # shape = [batch_size, dims_a, max_logit_length, max_label_length + 1, dims_b]

        horizontal_blank_grad_term = expand_many_dims(self.blank_logproba, axes=[1, 3]) + tf.reduce_logsumexp(ab_term, axis=3)
        # shape = [batch_size, dims_a, max_logit_length, dims_b]

        act = a[:, :, :, :, 1] + expand_many_dims(self.previous_label_token_log_proba, axes=[1, 4]) + b[:, :, :, :, 1]
        # shape = [batch_size, dim_a, max_logit_length, max_label_length + 1, dim_b]

        horizontal_non_blank_grad_term = self.select_from_act(act, self.preceded_label)
        # shape = [batch_size, dim_a, max_logit_length, num_tokens, dim_b]

        input_tensor = a + expand_many_dims(self.any_to_open_diagonal_step_log_proba, axes=[1, 5]) + tf.roll(b[:, :, :, :, 1:], shift=-1, axis=3)
        # shape = [batch_size, dim_a, max_logit_length, max_label_length + 1, states, dim_b]

        act = tf.reduce_logsumexp(input_tensor=input_tensor, axis=4)
        # shape = [batch_size, dim_a, max_logit_length, max_label_length + 1, dim_b]

        diagonal_non_blank_grad_term = self.select_from_act(act=act, label=self.label)
        # shape = [batch_size, dim_a, max_logit_length, num_tokens, dim_b]

        non_blank_grad_term = logsumexp(horizontal_non_blank_grad_term, diagonal_non_blank_grad_term)
        # shape = [batch_size, dim_a, max_logit_length, num_tokens, dim_b]

        blank_mask = self.blank_token_index == tf.range(self.num_tokens)
        # shape = [num_tokens]

        output = tf.where(
            condition=expand_many_dims(blank_mask, axes=[0, 1, 2, 4]),
            x=tf.expand_dims(horizontal_blank_grad_term, 3),
            y=non_blank_grad_term,
        )
        # shape = [batch, dim_a, max_logit_length, num_tokens, dim_b]
        output_shape = tf.concat(
            [
                tf.expand_dims(self.batch_size, axis=0),
                dims_a,
                tf.expand_dims(self.max_logit_length, axis=0),
                tf.expand_dims(self.num_tokens, axis=0),
                dims_b,
            ],
            axis=0,
        )
        output_reshaped = tf.reshape(output, shape=output_shape)
        # shape = [batch, DIMS_A, max_logit_length, num_tokens, DIMS_B]

        return output_reshaped
Transforms logarithmic transition probabilities a and b.

        Args:
            a:      shape = [batch, DIMS_A, max_logit_length, max_label_length + 1, state]
            b:      shape = [batch, max_logit_length, max_label_length + 1, state, DIMS_B]

        Returns:    shape = [batch, DIMS_A, max_logit_length, num_tokens, DIMS_B]
        
tensorflow_asr/losses/impl/ctc_tpu.py