def forward(
        ctx,
        p: torch.Tensor,
        q: torch.Tensor,
        shift_labels: Optional[torch.Tensor] = None,
        reduction: REDUCTION_LITERAL = "batchmean",
        ignore_index: int = -100,
    ) -> torch.Tensor:
        """A forward pass for the Total Variation Distance Loss.

        Args:
            ctx: Torch autograd context
            p (torch.Tensor): A tensor of shape (BT, V) containing the first distribution.
            q (torch.Tensor): A tensor of shape (BT, V) containing the second distribution.
            shift_labels (Optional[torch.Tensor]): A tensor of shape (BT,) containing the labels.
            reduction (REDUCTION_LITERAL, optional): The reduction method to be applied. Defaults to "batchmean".
            ignore_index (int, optional): The index to ignore during loss calculation. Defaults to -100.

        Returns:
            torch.Tensor: The computed Total Variation Distance Loss.
        """
        has_label = False
        if shift_labels is not None:
            assert shift_labels.shape == (p.shape[0],), (
                f"the shape of shift_labels must be (BT,). Got: {shift_labels.shape}"
            )
            shift_labels = shift_labels.contiguous()
            has_label = True

        loss, grads = tv_distance_forward_triton(p, q, shift_labels, reduction, ignore_index, has_label)
        ctx.save_for_backward(grads)
        return loss
A forward pass for the Total Variation Distance Loss.

        Args:
            ctx: Torch autograd context
            p (torch.Tensor): A tensor of shape (BT, V) containing the first distribution.
            q (torch.Tensor): A tensor of shape (BT, V) containing the second distribution.
            shift_labels (Optional[torch.Tensor]): A tensor of shape (BT,) containing the labels.
            reduction (REDUCTION_LITERAL, optional): The reduction method to be applied. Defaults to "batchmean".
            ignore_index (int, optional): The index to ignore during loss calculation. Defaults to -100.

        Returns:
            torch.Tensor: The computed Total Variation Distance Loss.
        
src/liger_kernel/ops/tvd.py