def gamma_step(
        self,
        previous_slice: tf.Tensor,
        i: tf.Tensor,
    ) -> tf.Tensor:
        """Args:
            previous_slice: tf.Tensor,
                            shape = [batch_size, max_logit_length + 1, max_label_length + 1, state,
                                max_label_length + 1, state]
            i:              tf.Tensor,
                            shape = [], 0 <= i < max_logit_length + 1

        Returns:            tf.Tensor,
                            shape = [batch_size, max_logit_length + 1, max_label_length + 1, state,
                                max_label_length + 1, state]
        """
        horizontal_step_states = expand_many_dims(self.horizontal_step_log_proba[:, i], axes=[1, 2, 3]) + tf.expand_dims(previous_slice, 5)
        # shape = [batch_size, max_logit_length + 1, max_label_length + 1, state,
        #          max_label_length + 1, next_state, previous_state]
        horizontal_step = tf.reduce_logsumexp(horizontal_step_states, axis=6)
        # shape = [batch_size, max_logit_length + 1, max_label_length + 1, state, max_label_length + 1, state]

        diagonal_step_log_proba = tf.reduce_logsumexp(
            expand_many_dims(self.any_to_open_diagonal_step_log_proba[:, i], axes=[1, 2, 3]) + previous_slice, axis=5
        )
        # shape = [batch_size, max_logit_length + 1, max_label_length + 1, state, max_label_length + 1]

        # We move by one token because it is a diagonal step
        moved_diagonal_step_log_proba = tf.roll(diagonal_step_log_proba, shift=1, axis=4)
        # shape = [batch_size, max_logit_length + 1, max_label_length + 1, state, max_label_length + 1]

        # Out state is always open:
        diagonal_step = tf.pad(
            tensor=tf.expand_dims(moved_diagonal_step_log_proba, 5),
            paddings=[[0, 0], [0, 0], [0, 0], [0, 0], [0, 0], [1, 0]],
            constant_values=-np.inf,
        )
        # shape = [batch_size, max_logit_length + 1, max_label_length + 1, state, max_label_length + 1, state]
        new_gamma_slice = logsumexp(
            x=horizontal_step,
            y=diagonal_step,
        )
        # shape = [batch_size, max_logit_length + 1, max_label_length + 1, state, max_label_length + 1, state]

        condition = tf.reshape(tf.range(self.max_logit_length_plus_one) <= i, shape=[1, -1, 1, 1, 1, 1])
        # shape = [1, max_logit_length + 1, 1, 1, 1, 1, 1]
        output_slice = tf.where(
            condition=condition,
            x=new_gamma_slice,
            y=self.diagonal_gamma,
        )
        # shape = [batch_size, max_logit_length + 1, max_label_length + 1, state, max_label_length + 1, state]

        return output_slice
Args:
            previous_slice: tf.Tensor,
                            shape = [batch_size, max_logit_length + 1, max_label_length + 1, state,
                                max_label_length + 1, state]
            i:              tf.Tensor,
                            shape = [], 0 <= i < max_logit_length + 1

        Returns:            tf.Tensor,
                            shape = [batch_size, max_logit_length + 1, max_label_length + 1, state,
                                max_label_length + 1, state]
        
tensorflow_asr/losses/impl/ctc_tpu.py