"""Mirrors ``tilelang_ops/asc/dflash.py`` (confidence_head_fused_tl, dflash_input_store_tl) for DSpark drafting:
the frontend pure-torch references with the kernels' argument lists."""

from typing import Any

import torch


def confidence_head_fused_tl(n_embed: int, markov_rank: int, has_markov_emb: bool) -> Any:
    from ..dflash import confidence_head_fused_ref

    def _confidence(hidden, markov_emb, w, out, temperature):  # noqa: ANN001
        confidence_head_fused_ref(hidden, w.view(-1, 1), temperature, markov_emb if has_markov_emb else None, out=out)

    return _confidence


def dflash_input_store_tl(n_embed: int, hc_mult: int) -> Any:
    def _store(hidden_states, out):  # noqa: ANN001
        h = hidden_states.float().mean(dim=-2) if hc_mult > 1 else hidden_states.reshape(-1, n_embed)
        out.copy_(h.to(out.dtype).reshape(out.shape))

    return _store