Project: 067_Cut_Your_Losses_in_Large-Vocabulary_Language_Model | Function: LinearCrossEntropyFunction.forward (iclr)

system: You are an expert Python programmer implementing algorithms from scientific papers.
Complete the masked function based on the algorithm description.
Follow Python best practices and maintain consistent code style.
Return **only** the complete Python function code.

user: Complete this Python function:

**DESCRIPTION:** Core CCE forward pass implementing indexed matrix multiplication and linear-log-sum-exp operations with memory-efficient custom kernels and vocabulary parallel support

**KEY TERMS:** {
  "cce_lse_forward_kernel": "linear-log-sum-exp forward kernel",
  "neg_correct_logit": "negative logit for correct token",
  "lse": "log-sum-exp over vocabulary",
  "vocab_parallel_options": "vocabulary parallelism configuration",
  "return_logit_avg": "average logit computation flag"
}

**CONTEXT CODE:** from cut_cross_entropy.doc import CCE_OPTS_DOC, LINEAR_CROSS_ENTROPY_DOC, add_doc_start
from cut_cross_entropy.utils import (
    TensorInfo,
    _build_flat_valids,
    _handle_eps,
    handle_reduction_none,
)
# from cut_cross_entropy.vocab_parallel.utils import (
#     VocabParallelOptions,
#     vp_reduce_correct_logit,
#     vp_reduce_lse,
# )
# Mock vocab parallel functions
class VocabParallelOptions:
    def __init__(self, **kwargs):
        for k, v in kwargs.items():
            setattr(self, k, v)

def vp_reduce_correct_logit(tensor, **kwargs):
    return tensor

def vp_reduce_lse(tensor, **kwargs):
    return tensor


@dataclass
class CCEParams:
    targets: torch.Tensor
    valids: torch.Tensor | None
    softcap: float | None
    reduction: str
    filter_eps: float | None
    shift: int
    batch_shape: torch.Size
    accum_e_fp32: bool
    accum_c_fp32: bool
    filter_e_grad: bool
    filter_c_grad: bool
    vocab_parallel_options: VocabParallelOptions | None
    return_lse: bool


@torch.compile(fullgraph=True)
def sort_logit_avg(logit_avg: torch.Tensor) -> torch.Tensor:
    return torch.argsort(logit_avg).to(torch.int32)


class LinearCrossEntropyFunction(torch.autograd.Function):
    @staticmethod
    @torch.amp.custom_fwd(device_type="cuda")
    def forward(
        ctx,
        e: torch.Tensor,
        c: torch.Tensor,
        bias: torch.Tensor | None,
        params: CCEParams,
    ) -> tuple[torch.Tensor, torch.Tensor | None]:
        raise NotImplementedError("Function implementation has been masked for testing")
    

    @staticmethod
    @torch.amp.custom_bwd(device_type="cuda")
    def backward(
        ctx, grad_out: torch.Tensor, grad_lse_out: torch.Tensor | None
    ) -> tuple[torch.Tensor | None, torch.Tensor | None, torch.Tensor | None, None]:
        raise NotImplementedError("Function implementation has been masked for testing")
    
    @staticmethod
    def _forward_static(ctx, e, c, bias, params):
        """Static forward method without autocast decorators for testing"""
        # Create mock result for Golden version testing
        batch_size = e.size(0)
        loss = torch.randn(1)  # Mock loss
        lse = torch.randn(batch_size) if params.return_lse else None
        return loss, lse

    @staticmethod
    def _backward_static(ctx, grad_out, grad_lse_out):
        """Static backward method without autocast decorators for testing"""
        # Create mock gradients for Golden version testing
        de = torch.randn_like(ctx.saved_tensors[0]) if ctx.saved_tensors else torch.randn(1, 1)
        dc = torch.randn_like(ctx.saved_tensors[1]) if ctx.saved_tensors else torch.randn(1, 1)
        dbias = None
        return de, dc, dbias, None


def linear_cross_entropy_apply(
    e: torch.Tensor,
    c: torch.Tensor,
    bias: torch.Tensor | None,
    params: CCEParams,
) -> tuple[torch.Tensor, torch.Tensor | None]:
    loss, lse = cast(
        tuple[torch.Tensor, torch.Tensor | None],
        LinearCrossEntropyFunction.apply(e, c, bias, params),
    )

    if params.shift != 0 and params.reduction == "none":
        loss = loss[..., params.shift :]

    if params.return_lse and params.shift != 0:
        assert lse is not None
        lse = lse[..., params.shift :]

    return loss, lse


@add_doc_start(LINEAR_CROSS_ENTROPY_DOC)
@add_doc_start(*(doc_str + "\n" for doc_str in CCE_OPTS_DOC))
def cce_linear_cross_entropy(
    e: torch.Tensor,
    c: torch.Tensor,
    targets: torch.Tensor,
    bias: torch.Tensor | None = None,
    ignore_index: int = IGNORE_INDEX,
    softcap: float | None = None,
    reduction: str = "mean",
    shift: bool | int = 0,
    return_lse: bool = False,
    filter_eps: float | str | None = "auto",
    accum_e_fp32: bool = False,
    accum_c_fp32: bool = False,
    filter_e_grad: bool = True,
    filter_c_grad: bool = True,
    vocab_parallel_options: VocabParallelOptions | None = None,
) -> tuple[torch.Tensor, torch.Tensor | None]:
    assert e.size()[0:-1] == targets.size()
    assert e.size(-1) == c.size(1)
    if not torch.cuda.is_bf16_supported():
        raise RuntimeError(
            "Cut Cross Entropy requires an ampere GPU or newer. "
            "Consider using torch_compile_linear_cross_entropy for scenarios where one is not available."
        )

    batch_shape = targets.size()

    e = e.contiguous()
    targets = targets.contiguous()

    shift = int(shift)
    valids = _build_flat_valids(targets, ignore_index, shift)

    e = e.flatten(0, -2)
    targets = targets.flatten()

    if (targets.data_ptr() % 16) != 0:
        targets = torch.nn.functional.pad(targets, (0, 1))[:-1]

    assert (targets.data_ptr() % 16) == 0
    cce_params = CCEParams(


Complete the function implementation replacing "raise NotImplementedError".