def hessian(self) -> tf.Tensor:
        """Calculates Hessian of loss w.r.t. input logits.

        Returns: tf.Tensor, shape = [batch_size, max_logit_length, num_tokens, max_logit_length, num_tokens]
        """
        alpha_gamma_term = self.combine_transition_probabilities(a=self.alpha[:, :-1], b=self.gamma[:, 1:])
        # shape = [batch_size, max_logit_length, num_tokens, max_logit_length + 1, max_label_length + 1]
        alpha_gamma_beta_term = self.combine_transition_probabilities(a=alpha_gamma_term[:, :, :, :-1], b=self.beta[:, 1:])
        # shape = [batch_size, max_logit_length, num_tokens, max_logit_length, num_tokens]
        alpha_gamma_beta_loss_term = expand_many_dims(self.loss, axes=[1, 2, 3, 4]) + alpha_gamma_beta_term
        # shape = [batch_size, max_logit_length, num_tokens]
        logit_length_x_num_tokens = self.max_logit_length * self.num_tokens
        first_term = tf.reshape(
            tf.linalg.set_diag(
                input=tf.reshape(tensor=alpha_gamma_beta_loss_term, shape=[self.batch_size, logit_length_x_num_tokens, logit_length_x_num_tokens]),
                diagonal=tf.reshape(tensor=self.logarithmic_logproba_gradient, shape=[self.batch_size, logit_length_x_num_tokens]),
            ),
            shape=tf.shape(alpha_gamma_beta_term),
        )

        mask = expand_many_dims(input=tf.linalg.band_part(tf.ones(shape=[self.max_logit_length] * 2, dtype=tf.bool), 0, -1), axes=[0, 2, 4])
        symmetrized_first_term = tf.where(
            condition=mask,
            x=first_term,
            y=tf.transpose(first_term, [0, 3, 4, 1, 2]),
        )
        # shape = [batch_size, max_logit_length, num_tokens, max_logit_length, num_tokens]
        hessian = -tf.exp(symmetrized_first_term) + expand_many_dims(self.gradient, [3, 4]) * expand_many_dims(self.gradient, [1, 2])
        # shape = [batch_size, max_logit_length, num_tokens, max_logit_length, num_tokens]

        # Filter out samples with infinite loss
        hessian = tf.where(
            condition=expand_many_dims(self.loss == inf, [1, 2, 3, 4]),
            x=tf.zeros(shape=[1, 1, 1, 1, 1]),
            y=hessian,
        )
        # shape = [batch_size, max_logit_length, num_tokens, max_logit_length, num_tokens]

        # Filter out logits that beyond logits length
        hessian = tf.where(condition=expand_many_dims(self.logit_length_mask, axes=[2, 3, 4]), x=hessian, y=0.0)
        hessian = tf.where(condition=expand_many_dims(self.logit_length_mask, axes=[1, 2, 4]), x=hessian, y=0.0)

        return hessian
Calculates Hessian of loss w.r.t. input logits.

        Returns: tf.Tensor, shape = [batch_size, max_logit_length, num_tokens, max_logit_length, num_tokens]
        
tensorflow_asr/losses/impl/ctc_tpu.py