def update(
        self,
        labels: torch.Tensor,
        logits: torch.Tensor,
        sequence_id: Optional[torch.Tensor],
        loss_fn: Any,
    ) -> None:
        """Updates the internal state with results from a new batch.

        Args:
            labels (torch.Tensor): A Tensor of ground-truth values to compare against.
            logits (torch.Tensor): A Tensor of labels.
            sequence_id (torch.Tensor | None): The sequence ids for tokens.
            loss_fn (Any): The cross entropy loss to use.
        """
        valid_labels_mask = torch.where(
            labels != self.ignore_index,
            torch.ones_like(labels),
            torch.zeros_like(labels),
        )
        bsz, seq_len = labels.shape
        loss = loss_fn(logits.view(bsz * seq_len, -1), labels.view(-1))
        loss = loss.view(bsz, seq_len)
        perplexity = torch.exp(loss)

        if self.sum_loss.numel() == 0:
            self.sum_loss = torch.zeros(  # type: ignore
                seq_len,
                device=loss.device,
                dtype=loss.dtype,
            )
            self.sum_perplexity = torch.zeros(  # type: ignore
                seq_len,
                device=loss.device,
                dtype=loss.dtype,
            )
            self.sum_length = torch.zeros(  # type: ignore
                seq_len,
                device=loss.device,
                dtype=torch.long,
            )
            self.sum_loss_seq_id = torch.zeros(  # type: ignore
                seq_len,
                device=loss.device,
                dtype=loss.dtype,
            )
            self.sum_perplexity_seq_id = torch.zeros(  # type: ignore
                seq_len,
                device=loss.device,
                dtype=loss.dtype,
            )
            self.sum_length_seq_id = torch.zeros(  # type: ignore
                seq_len,
                device=loss.device,
                dtype=torch.long,
            )

        self.sum_loss += torch.sum(loss, dim=(0))
        self.sum_perplexity += torch.sum(perplexity, dim=(0))
        self.sum_length += valid_labels_mask.sum(dim=0)

        if sequence_id is not None:
            seq_id_mask = (sequence_id != -1)
            sequence_id = torch.where(seq_id_mask, sequence_id, 0)
            seq_id_expanded = torch.nn.functional.one_hot(
                sequence_id,
            )
            seq_id_expanded = torch.where(
                torch.unsqueeze(seq_id_mask, dim=-1),
                seq_id_expanded,
                0,
            )
            seq_id_expanded = seq_id_expanded.transpose(-1, -2)
            seq_lens = seq_id_expanded.sum(dim=-1)
            max_num_seq = seq_lens.shape[1]
            seq_tok_ids = torch.arange(seq_len, device=sequence_id.device)[
                None, None, :].expand(bsz, max_num_seq, -1)
            mask = seq_tok_ids < seq_lens[:, :, None]
            seq_len_offsets = torch.nn.functional.pad(
                seq_lens.cumsum(dim=1)[:, :-1],
                (1, 0),
                value=0,
            )
            seq_tok_ids = seq_tok_ids + seq_len_offsets[:, :, None]
            seq_tok_ids = torch.where(
                mask,
                seq_tok_ids,
                torch.zeros_like(seq_tok_ids),
            )

            loss = loss[:, None, :].expand(-1, max_num_seq, -1)
            perplexity = perplexity[:, None, :].expand(-1, max_num_seq, -1)
            valid_labels_mask = valid_labels_mask[:, None, :].expand(
                -1,
                max_num_seq,
                -1,
            )
            loss = torch.where(
                mask,
                torch.gather(input=loss, dim=2, index=seq_tok_ids),
                torch.zeros_like(loss),
            )
            perplexity = torch.where(
                mask,
                torch.gather(input=perplexity, dim=2, index=seq_tok_ids),
                torch.zeros_like(perplexity),
            )
            mask = torch.where(
                mask,
                torch.gather(input=valid_labels_mask, dim=2, index=seq_tok_ids),
                torch.zeros_like(valid_labels_mask),
            )

            self.sum_loss_seq_id += torch.sum(loss, dim=(0, 1))
            self.sum_perplexity_seq_id += torch.sum(perplexity, dim=(0, 1))
            self.sum_length_seq_id += torch.sum(mask, dim=(0, 1))
Updates the internal state with results from a new batch.

        Args:
            labels (torch.Tensor): A Tensor of ground-truth values to compare against.
            logits (torch.Tensor): A Tensor of labels.
            sequence_id (torch.Tensor | None): The sequence ids for tokens.
            loss_fn (Any): The cross entropy loss to use.
        
llmfoundry/callbacks/loss_perp_v_len_callback.py