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

        Returns:            shape = [batch_size, max_label_length + 1, state]
        """
        temp = self.horizontal_step_log_proba[:, i] + tf.expand_dims(previous_slice, 2)
        # shape = [batch_size, max_label_length + 1, next_state, previous_state]
        horizontal_step = tf.reduce_logsumexp(temp, 3)
        # shape = [batch_size, max_label_length + 1, state]
        diagonal_step_log_proba = tf.reduce_logsumexp(self.any_to_open_diagonal_step_log_proba[:, i] + previous_slice, 2)
        # shape = [batch_size, 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=1)
        # shape = [batch_size, max_label_length + 1]

        # Out state is always open:
        diagonal_step = tf.pad(tensor=tf.expand_dims(moved_diagonal_step_log_proba, 2), paddings=[[0, 0], [0, 0], [1, 0]], constant_values=-np.inf)
        # shape = [batch_size, max_label_length + 1, state]
        new_alpha_slice = logsumexp(
            x=horizontal_step,
            y=diagonal_step,
        )
        # shape = [batch_size, max_label_length + 1, state]
        return new_alpha_slice
Args:
            previous_slice: shape = [batch_size, max_label_length + 1, state]
            i:

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