mlx-bun library API
    Preparing search index...
    • B=1 FUSED linear cross-entropy response-only mean log-prob → [1]. Same result as branchLogpMeanB1 but the LM head is a single CustomVjp with an ANALYTIC backward — the Liger FusedLinearCrossEntropy structure ported to MLX:

      • FORWARD token-chunks the M response rows; per chunk materializes only [Cc, V] logits (= h_c @ Wᵀ via quantizedMatmul, + Gemma softcap), reduces to logp = target_logit − logsumexp, and DROPS the logits. Returns Σ logp (the 1/M mean is applied outside, a cheap linear scale).
      • BACKWARD recomputes [Cc, V] per chunk and forms the gradient in CLOSED FORM — ∂(Σlogp)/∂logit = onehot_target − softmax (× sech² through the softcap), scaled by the incoming cotangent — instead of autograd-ing through logsumexp+gather (what the Checkpoint token-chunked path does). dh_c = grad_logits_c @ dequant(W) uses quantizedMatmul(…, transpose=false), which IS mlx's own x-vjp of the transpose=true head matmul (the path LoRA training already differentiates through), so it contracts the vocab axis correctly without dequantizing the full [V, hidden] head (1.3 GB at 262k). Peak large-vocab memory is bounded to [chunkSize, V] in BOTH directions, the full [M, V] logits never materialize, and the head backward needs no autograd graph. The CustomVjp recomputes in the backward (it must outlive the enclosing value_and_grad), so it — and h — are pushed into sink, disposed by the caller after the step's grads are eval'd. The head is NOT a LoRA target, so only dh flows back (seeding the layer-stack LoRA backward); a head-LoRA ∇A/∇B fold-in is a follow-on. Exact within the established bf16 class.

      Parameters

      • model: RuntimeModel
      • ids: number[]
      • mask: number[]
      • validLen: number
      • chunkSize: number
      • sink: { dispose(): void }[]
      • vocabBlock: number = 0
      • flash: FlashCce | null = null

      Returns MlxArray