legume-numeric 0.8.11

//! Embedded topic decoder with **negative-binomial masked imputation**.
//!
//! The training head of the masked-imputation topic model. Same ETM
//! factorization as a dense embedded topic decoder
//! (`β = softmax_d(α·ρᵀ)`, `ρ` shared with the encoder), but instead of a
//! multinomial reconstruction it scores a **negative-binomial** likelihood on
//! the **held-out (masked)** genes only:
//!
//! ```text
//! μ_gn = residual_gn · ℓ_n · (θ_n · β)_g          (residual = per-cell μ_residual offset)
//! x_gn ~ NB(μ_gn, φ_g)                            (φ_g per-gene dispersion)
//! llik = Σ_{g ∈ masked} log NB(x_gn | μ_gn, φ_g)
//! ```
//!
//! No KL / no variational posterior is involved — `θ` is the encoder's point
//! estimate — so the objective has no posterior-collapse pressure. Modelling
//! observed counts as `NB(residual · ℓ · θβ)` keeps `β` batch-free (the
//! per-cell `residual` absorbs the batch effect, matching the collapse model
//! `E[y] = μ_residual · μ_adjusted`).

use crate::candle::batched_dot::batched_matvec;
use crate::candle::decoder::coarsening_map::CoarseningMap;
use crate::candle::fast_index::gather_rows;
use crate::candle::loss::nb_log_likelihood_elem;
use candle_core::{Result, Tensor};
use candle_nn::{ops, VarBuilder};

/// Per-cell minibatch target for [`EmbeddedNbTopicDecoder::impute_masked_nb`].
/// All tensors are at the cell's top-K positions, in `indices` order.
pub struct MaskedNbTarget<'a> {
    /// `[N, K]` u32 per-cell gene ids (the cell's top-K).
    pub indices: &'a Tensor,
    /// `[N, K]` per-cell μ_residual at `indices` (batch offset); `None` ⇒ no
    /// batch offset (factor 1).
    pub residual: Option<&'a Tensor>,
    /// `[N, K]` observed counts at `indices` (NB target).
    pub values: &'a Tensor,
    /// `[N, 1]` per-cell library size.
    pub lib: &'a Tensor,
    /// `[N, K]` 1 = masked (scored), 0 = visible.
    pub mask: &'a Tensor,
}

/// Minibatch target for the **dense** masked heads: the row laid out over the
/// whole gene axis, plus the ids of the genes it is scored at.
///
/// This is the canonical masked-training shape: the encoder reads a bounded
/// context, and the decoder is scored over a prediction space that does not
/// depend on that budget, so a gene the encoder never saw — including one with
/// a zero count — still teaches the dictionary.
///
/// The scored set arrives as **ids**, not as an `[N, D]` indicator. Both say
/// the same thing, but a mask says it only after the elementwise likelihood has
/// already been computed over every gene and then multiplied to zero at the
/// visible ones — the majority of the axis. The ids let the head evaluate the
/// likelihood where it is scored and nowhere else. The dense parts of the rate
/// (the `[N, K] × [K, D]` product and the partition over all `D`) are
/// untouched: the partition is a normalisation over the whole gene axis and
/// restricting it would change the model, not just its cost.
pub struct MaskedDenseTarget<'a> {
    /// `[N, D]` observed counts over the full gene axis.
    pub values: &'a Tensor,
    /// `[N, D]` per-row μ_residual (batch offset); `None` ⇒ no offset. Training
    /// scores batch-free target rows and passes `None`; cell-level scoring,
    /// where the counts are batch-mixed, passes the per-cell offset.
    pub residual: Option<&'a Tensor>,
    /// `[N, 1]` per-row library size over the full row.
    pub lib: &'a Tensor,
    /// `[N, d_h]` u32 ids of the scored (hidden) genes, ascending within a row.
    pub hidden_ids: &'a Tensor,
    /// `[N, d_h]` 1 on a scored slot, 0 on a pad — needed only when the rows
    /// hide different counts and the block was padded to the widest of them
    /// (see [`crate::candle::data::masked_dense::DenseMaskedMinibatch::hidden_ids`]).
    /// `None` when every row hides exactly `d_h`: then there is nothing to
    /// multiply by.
    pub hidden_weight: Option<&'a Tensor>,
}

/// Minibatch target for the **module-collapsed** head: the row's counts
/// summed into coarse features, less what the encoder's context saw.
///
/// The module head predicts each module's UNSEEN share. The visible genes'
/// counts are subtracted from the module total (both from the target row, so
/// the difference is exact), and the module's rate is scaled by the pinned
/// share of its mass the context did not see. A module the context saw
/// completely has nothing left to predict and is not scored. Under the
/// identity map this is the dense gene head term for term.
pub struct ModuleTarget<'a> {
    /// `[N, M]` module totals of the target row.
    pub values: &'a Tensor,
    /// `[N, M]` target counts of the visible context genes, summed per module.
    pub visible_counts: &'a Tensor,
    /// `[N, M]` `Σ_{g visible in m} π_{g|m}`: the module's share the context saw.
    pub visible_share: &'a Tensor,
    /// `[N, M]` per-row multiplicative batch offset per module (the per-gene
    /// offset averaged by share); `None` ⇒ none. Training scores batch-free
    /// rows and passes `None`; cell-level scoring passes the cell's offset.
    pub residual: Option<&'a Tensor>,
    /// `[N, 1]` per-row library size over the full row.
    pub lib: &'a Tensor,
}

/// Minibatch target for the gene-level **query head**: a sampled set of
/// genes per row, each read at its module's rate times its share times the
/// query decoder's residual.
pub struct QueryTarget<'a> {
    /// `[N, Q]` u32 query genes (gene 0 on pads).
    pub gene_ids: &'a Tensor,
    /// `[N, Q]` target counts at the query genes.
    pub values: &'a Tensor,
    /// `[N, Q]` 1 on a real query, 0 on a pad.
    pub weight: &'a Tensor,
    /// `[N, Q]` the query decoder's log-residual `r_g`.
    pub log_residual: &'a Tensor,
    /// `[N, 1]` per-row library size over the full row.
    pub lib: &'a Tensor,
}

/// NB embedded-topic decoder for masked imputation.
pub struct EmbeddedNbTopicDecoder {
    /// Rows of ρ: the gene axis.
    n_features: usize,
    /// The decoder's output width: coarse features under a coarsening, every
    /// feature under
    /// the identity.
    n_obs: usize,
    n_topics: usize,
    /// Gene → module map; identity when the decoder scores genes.
    coarsening: CoarseningMap,
    /// `α [K, H]` topic embeddings (learnable, decoder scope).
    topic_embeddings: Tensor,
    /// The feature side, shared with the encoder (ETM tying) so gradients from
    /// either path land on the same parameters. Held as a handle rather than a
    /// table: under modules the rows are composed, and a copy would freeze
    /// them at construction.
    features: std::sync::Arc<crate::candle::feature_embedding::FeatureEmbedding>,
    /// `log φ_g [1, D]` per-gene NB inverse dispersion (learnable).
    log_phi_1d: Tensor,
    /// `log π_g [1, D]` per-gene log-background, added to every topic's logits
    /// and **pinned** at the data's gene marginal (see [`pin_background`]).
    ///
    /// Centering `α` over the topic axis makes each gene's total log-mass a
    /// conserved quantity, so a gene abundant in *every* cell has nowhere to
    /// put that abundance: the only escapes are a permanently dead
    /// "background" topic or diffuse `θ`. `log π_g` is that home. It must stay
    /// frozen — a learnable per-gene bias shifts all `K` topics equally on a
    /// gene, which is exactly the shared direction centering removes, and the
    /// optimizer would reinstate it through this parameter.
    log_pi_1d: Tensor,
}

/// Name of the pinned background var inside a decoder's scope.
pub const BACKGROUND_VAR: &str = "log_pi";

/// Pin a decoder's background at `log_pi_1d` `[1, D]` by writing the var named
/// `{prefix}.log_pi` in `varmap`. The decoder holds a handle into that var, so
/// the value takes effect immediately and round-trips through `VarMap::save`;
/// the trainer excludes every `*.log_pi` var from the optimizer.
pub fn pin_background(varmap: &candle_nn::VarMap, prefix: &str, log_pi_1d: &Tensor) -> Result<()> {
    let name = format!("{prefix}.{BACKGROUND_VAR}");
    let tbl = varmap.data().lock().unwrap();
    let var = tbl
        .get(&name)
        .ok_or_else(|| candle_core::Error::Msg(format!("no var `{name}` to pin")))?;
    var.set(&log_pi_1d.to_device(var.device())?.to_dtype(var.dtype())?)
}

/// `log π_g` from a per-gene mean expression `mean_d` (any non-negative scale):
/// the normalized marginal, floored so an unobserved gene keeps a finite
/// background.
pub fn log_background_from_mean(mean_d: &[f32], device: &candle_core::Device) -> Result<Tensor> {
    let total: f64 = mean_d.iter().map(|&m| f64::from(m.max(0.0))).sum();
    let d = mean_d.len().max(1) as f64;
    let floor = 1e-3 / d;
    let log_pi: Vec<f32> = mean_d
        .iter()
        .map(|&m| {
            let p = if total > 0.0 {
                f64::from(m.max(0.0)) / total
            } else {
                1.0 / d
            };
            p.max(floor).ln() as f32
        })
        .collect();
    Tensor::from_vec(log_pi, (1, mean_d.len()), device)
}

impl EmbeddedNbTopicDecoder {
    /// Construct with a shared feature-embedding handle
    /// (`encoder.features_shared()`). `α` is Kaiming-init in `vs`;
    /// `log φ` starts at ln(2) ≈ 0.69 (moderate dispersion).
    pub fn new(
        n_topics: usize,
        features: std::sync::Arc<crate::candle::feature_embedding::FeatureEmbedding>,
        vs: VarBuilder,
    ) -> Result<Self> {
        let identity = CoarseningMap::identity(features.n_features(), features.device())?;
        Self::new_with_coarsening(n_topics, features, identity, vs)
    }

    /// A decoder whose dense output axis is the map's coarse features: `φ` and the
    /// pinned background live at `[1, M]`, the logits are `(α − ᾱ)·ρ̄ᵀ + log π_m`
    /// with `ρ̄` the within-module mean of ρ, and any gene is scored at its
    /// module's rate times its pinned share. Same var names as [`Self::new`].
    pub fn new_with_coarsening(
        n_topics: usize,
        features: std::sync::Arc<crate::candle::feature_embedding::FeatureEmbedding>,
        coarsening: CoarseningMap,
        vs: VarBuilder,
    ) -> Result<Self> {
        let n_features = features.n_features();
        let embedding_dim = features.embedding_dim();
        if coarsening.n_fine() != n_features {
            candle_core::bail!(
                "EmbeddedNbTopicDecoder: the coarsening covers {} features but ρ has {n_features}",
                coarsening.n_fine()
            );
        }
        let n_obs = coarsening.n_coarse();

        let init_ws = candle_nn::init::DEFAULT_KAIMING_NORMAL;
        let topic_embeddings =
            vs.get_with_hints((n_topics, embedding_dim), "topic.embeddings", init_ws)?;
        let log_phi_1d = vs.get_with_hints((1, n_obs), "log_phi", candle_nn::Init::Const(0.693))?;
        // Uniform until pinned: a constant is a no-op under the output-axis
        // log_softmax, so an unpinned decoder is exactly background-free.
        let log_pi_1d = vs.get_with_hints(
            (1, n_obs),
            BACKGROUND_VAR,
            candle_nn::Init::Const(-(n_obs as f64).ln()),
        )?;

        Ok(Self {
            n_features,
            n_obs,
            n_topics,
            coarsening,
            topic_embeddings,
            features,
            log_phi_1d,
            log_pi_1d,
        })
    }

    /// The feature → coarse-feature map (identity when the decoder scores
    /// every feature).
    pub fn coarsening(&self) -> &CoarseningMap {
        &self.coarsening
    }
    pub fn log_phi(&self) -> &Tensor {
        &self.log_phi_1d
    }
    /// Pinned per-gene log-background `[1, D]`.
    pub fn log_background(&self) -> &Tensor {
        &self.log_pi_1d
    }
    pub fn phi(&self) -> Result<Tensor> {
        self.log_phi_1d.exp()
    }
    /// The dense output width: modules under a module map, genes otherwise.
    pub fn dim_obs(&self) -> usize {
        self.n_obs
    }

    /// The gene axis (rows of ρ).
    pub fn n_features(&self) -> usize {
        self.n_features
    }
    pub fn dim_latent(&self) -> usize {
        self.n_topics
    }

    /// Full `[K, D]` pre-softmax logits `(α - ᾱ) · ρᵀ + log π_g`.
    ///
    /// `α` is centered over the topic axis first. The raw loading `α · ρᵀ` is
    /// dominated by a shared "abundance" direction (the mean archetype `ᾱ`) that
    /// ranks the same genes top in *every* topic; because each row is normalized
    /// independently, raising all `K` topics on those genes is a direction the
    /// whole dictionary descends at once, and nothing opposes it — the
    /// off-diagonal response of one topic to another is exactly zero.
    ///
    /// Subtracting `ᾱ` is a projection, so `Σ_k (α_k - ᾱ)·ρ_g = 0` holds for
    /// every gene at every step and for all parameter values: each gene's total
    /// log-mass becomes a conserved quantity and one topic can only gain on a
    /// gene at another's expense (`∂/∂S_jg` of topic `k` is `δ_jk - 1/K`, which
    /// is strictly negative off-diagonal). There is no coefficient to balance
    /// and no way to switch it off.
    ///
    /// This is the single chokepoint for the `[K, D]` logits — the dictionary,
    /// the log-partition and the training likelihood all read it — so the
    /// trained model and the dictionary written to disk stay consistent.
    pub fn full_logits_kd(&self) -> Result<Tensor> {
        // Under a module map the table is the within-module mean of ρ, so ρ
        // still trains through the dense head — each gene at 1/|m| of its
        // module's gradient.
        // Read the feature side live, and read it through the map: the
        // coarsening is an unweighted row mean, so coarsening the membership
        // and then composing gives the same table as the other order. Under a
        // real coarsening that is `[d, M]` work rather than `[D, H]`; under the
        // identity map, which is the default, it is the full table either way.
        let table = self
            .features
            .map_rows_linear(|rows| self.coarsening.coarsen_mean_dh(rows))?;
        self.centered_topic_embeddings()?
            .matmul(&table.t()?)?
            .broadcast_add(&self.log_pi_1d)
    }

    /// `α - ᾱ` `[K, H]`: the topic embeddings with the mean archetype removed.
    /// Every consumer of `α` — the full dictionary logits *and* the per-cell
    /// rate at the sampled genes — must go through this, because the
    /// log-partition is taken over the centered logits and the numerator has to
    /// be the same quantity or `β` stops summing to one.
    fn centered_topic_embeddings(&self) -> Result<Tensor> {
        let alpha_mean_1h = self.topic_embeddings.mean_keepdim(0)?;
        self.topic_embeddings.broadcast_sub(&alpha_mean_1h)
    }

    /// Full `[D, K]` log-β = `log_softmax_d((α - ᾱ)·ρᵀ + log π)` — for dictionary output.
    pub fn get_dictionary(&self) -> Result<Tensor> {
        let logits_kd = self.full_logits_kd()?;
        let log_beta_kd = ops::log_softmax(&logits_kd, logits_kd.rank() - 1)?;
        log_beta_kd.transpose(0, 1)?.contiguous()
    }

    /// Per-topic log-partition `log Z_k = logsumexp_d(logits_kd)` as `[1, 1, K]`,
    /// from precomputed `[K, D]` logits (see [`Self::full_logits_kd`]). The
    /// `[K, D]` product is the dominant decoder cost, so the caller computes it
    /// once per minibatch and shares it between this partition and the
    /// anchor-prior CE rather than recomputing it inside each.
    pub fn log_partition_from_logits(full_kd: &Tensor) -> Result<Tensor> {
        let k = full_kd.dim(0)?;
        Self::log_partition_k1(full_kd)?.reshape((1, 1, k))
    }

    /// Per-topic log-partition as `[K, 1]` — the natural shape for the dense
    /// `[K, D]` normalization. [`Self::log_partition_from_logits`] is the same
    /// quantity reshaped for the indexed head's `[N, K, T]` broadcast.
    pub fn log_partition_k1(full_kd: &Tensor) -> Result<Tensor> {
        let m = full_kd.max_keepdim(1)?; // [K, 1]
        let lse = (full_kd.broadcast_sub(&m)?.exp()?.sum_keepdim(1)? + 1e-20)?.log()?; // [K,1]
        lse + m
    }

    /// Per-row mixture rate `p_ng = Σ_t θ_nt · β_{t,g}` over the **whole** gene
    /// axis, `[N, D]`, so `Σ_g p_ng = 1` exactly.
    ///
    /// One `[N,K] × [K,D]` gemm against the normalized dictionary. The indexed
    /// sibling ([`Self::mixture_rate_nk`]) instead gathers `[N·K, T]` logits,
    /// which is the cheaper shape only while the scored set is a small
    /// per-row context.
    pub fn mixture_rate_nd(&self, log_theta_nk: &Tensor, full_kd: &Tensor) -> Result<Tensor> {
        let logz_k1 = Self::log_partition_k1(full_kd)?; // [K, 1]
        let beta_kd = full_kd.broadcast_sub(&logz_k1)?.exp()?; // [K, D]
        log_theta_nk.exp()?.matmul(&beta_kd) // [N, D]
    }

    /// Per-cell mixture rate `p_nk = Σ_t θ_nt · β_{t,g}` at the cell's top-K
    /// genes, with β normalized over the full vocab (so `Σ_g p_g = 1`). The
    /// shared core of both masked-impute heads; the NB and multinomial
    /// likelihoods differ only in how they score this rate. `[N, K]`.
    ///
    /// The logits are **gathered** from the caller's `full_kd` (see
    /// [`Self::full_logits_kd`]) rather than recomputed, so the numerator and
    /// the partition are one quantity by construction.
    pub(crate) fn mixture_rate_nk(
        &self,
        log_theta_nk: &Tensor,
        indices: &Tensor,
        full_kd: &Tensor,
    ) -> Result<Tensor> {
        let n = indices.dim(0)?;
        let k = indices.dim(1)?;
        let t = self.n_topics;

        let theta_nt = log_theta_nk.exp()?; // [N, T]
                                            // A gene is scored at its module's rate times its pinned share; under
                                            // the identity map both lookups are the gene itself and share one.
        let flat = self.coarsening.groups_of(indices)?.flatten_all()?; // [N*K] module ids

        let logz_11k = Self::log_partition_from_logits(full_kd)?; // [1, 1, T]
        let logits = gather_rows(&full_kd.t()?.contiguous()?, &flat)? // [M, T] → [N*K, T]
            .reshape((n, k, t))?; // [N, K, T]
        let beta_nkt = logits.broadcast_sub(&logz_11k)?.exp()?; // [N, K, T]

        // Mixture rate `Σ_t β·θ` as a gemm — see `legume_numeric::candle::batched_dot`.
        let rate_nk = batched_matvec(&beta_nkt, &theta_nt)?; // [N, K]
        if self.coarsening.is_identity() {
            return Ok(rate_nk);
        }
        rate_nk.mul(&self.coarsening.log_share_at(indices)?.exp()?)
    }

    /// Masked NB imputation log-likelihood, summed over masked positions →
    /// `[N]`.
    ///
    /// * `log_theta_nk` — `[N, K_topics]` encoder log-proportions.
    /// * `target` — the per-cell minibatch target (see [`MaskedNbTarget`]).
    /// * `full_kd` — `[K, D]` logits from [`Self::full_logits_kd`]
    ///   (caller-hoisted: it is the dominant decoder cost and the anchor-prior
    ///   CE shares it).
    pub fn impute_masked_nb(
        &self,
        log_theta_nk: &Tensor,
        target: &MaskedNbTarget<'_>,
        full_kd: &Tensor,
    ) -> Result<Tensor> {
        let MaskedNbTarget {
            indices,
            residual: residual_nk,
            values: values_nk,
            lib: lib_n1,
            mask: mask_nk,
        } = *target;

        let (n, k) = (indices.dim(0)?, indices.dim(1)?);
        let theta_beta_nk = self.mixture_rate_nk(log_theta_nk, indices, full_kd)?; // [N, K]

        // φ at the cell's genes: per module under a module map.
        let flat = self.coarsening.groups_of(indices)?.flatten_all()?; // [N*K]
        let log_phi_nk = gather_rows(&self.log_phi_1d.squeeze(0)?, &flat)?.reshape((n, k))?; // [N, K]

        nb_score(
            values_nk,
            &theta_beta_nk,
            residual_nk,
            lib_n1,
            &log_phi_nk,
            Some(mask_nk),
        )
    }

    /// Masked **multinomial** (categorical) imputation log-likelihood, summed
    /// over masked positions → `[N]`. The MLM-faithful sibling of
    /// [`Self::impute_masked_nb`]: it reuses the identical mixture rate
    /// `p_g = Σ_t θ_t · β_{t,g}` (β normalized over the full vocab, so
    /// `Σ_g p_g = 1`) but scores it as full-vocab categorical cross-entropy
    /// `Σ_{g∈mask} y_g · log p_g` — exactly BERT's MLM loss — instead of a
    /// per-gene NB. Depth-invariant: no library-size, no dispersion `φ`, no
    /// batch `residual` (those shape the NB *counts*; the multinomial models
    /// only relative composition). Sharing `p_g` with the NB head makes an
    /// ELBO-vs-masked comparison differ *only* in the objective, not the
    /// likelihood family.
    pub fn impute_masked_multinomial(
        &self,
        log_theta_nk: &Tensor,
        target: &MaskedNbTarget<'_>,
        full_kd: &Tensor,
    ) -> Result<Tensor> {
        let p_nk = self.mixture_rate_nk(log_theta_nk, target.indices, full_kd)?; // [N, K]
                                                                                 // Categorical cross-entropy at masked positions: Σ y_g · log p_g.
        multinomial_score(target.values, &p_nk, Some(target.mask))
    }

    ////////////////////////////////////////////////////
    // Dense heads — scored over the whole gene axis   //
    ////////////////////////////////////////////////////

    /// Masked NB imputation log-likelihood at the row's hidden genes → `[N]`.
    ///
    /// The dense sibling of [`Self::impute_masked_nb`]: same likelihood, same
    /// dictionary, but the scored positions come from the caller's hidden ids
    /// instead of a per-row context, so a gene with a zero count is a
    /// first-class observation rather than one the top-K dropped.
    ///
    /// The rate is formed over the WHOLE axis — one `[N, K] × [K, D]` product
    /// against `β` normalised over all `D` genes — and only then gathered to
    /// `[N, d_h]`. Restricting the product would build an `[N, d_h, K]` block,
    /// larger than it saves; restricting the partition would renormalise `β`
    /// over the scored genes and change the model.
    pub fn impute_dense_nb(
        &self,
        log_theta_nk: &Tensor,
        target: &MaskedDenseTarget<'_>,
        full_kd: &Tensor,
    ) -> Result<Tensor> {
        let ids = target.hidden_ids;
        let (n, dh) = ids.dims2()?;
        let rate_h = self
            .mixture_rate_nd(log_theta_nk, full_kd)?
            .gather(ids, 1)?;
        let values_h = target.values.contiguous()?.gather(ids, 1)?;
        let residual_h = target
            .residual
            .map(|r| r.contiguous()?.gather(ids, 1))
            .transpose()?;
        let log_phi_h =
            gather_rows(&self.log_phi_1d.squeeze(0)?, &ids.flatten_all()?)?.reshape((n, dh))?;
        nb_score(
            &values_h,
            &rate_h,
            residual_h.as_ref(),
            target.lib,
            &log_phi_h,
            target.hidden_weight,
        )
    }

    /// Masked multinomial imputation log-likelihood at the row's hidden genes →
    /// `[N]`. The dense sibling of [`Self::impute_masked_multinomial`].
    pub fn impute_dense_multinomial(
        &self,
        log_theta_nk: &Tensor,
        target: &MaskedDenseTarget<'_>,
        full_kd: &Tensor,
    ) -> Result<Tensor> {
        let ids = target.hidden_ids;
        let rate_h = self
            .mixture_rate_nd(log_theta_nk, full_kd)?
            .gather(ids, 1)?;
        let values_h = target.values.contiguous()?.gather(ids, 1)?;
        multinomial_score(&values_h, &rate_h, target.hidden_weight)
    }

    ////////////////////////////////////////////////////////
    // Module-collapsed head and the gene-level query head //
    ////////////////////////////////////////////////////////

    /// NB log-likelihood of each module's unseen counts under the module's
    /// rate scaled by its unseen share → `(llik [N], scored units [N])`.
    ///
    /// `μ_nm = ℓ_n · (θβ)_m · (1 − visible_share_nm)`; a module whose share
    /// the context saw completely is not scored. Under the identity map this
    /// equals [`Self::impute_dense_nb`] on the mask.
    pub fn score_unseen_modules_nb(
        &self,
        log_theta_nk: &Tensor,
        target: &ModuleTarget<'_>,
        full_km: &Tensor,
    ) -> Result<(Tensor, Tensor)> {
        let (rate_nm, unseen, share, scored) = self.unseen_parts(log_theta_nk, target, full_km)?;
        let mu = rate_nm.mul(&share)?.broadcast_mul(target.lib)?;
        let mu = match target.residual {
            Some(r) => mu.mul(r)?,
            None => mu,
        };
        let log_phi = self.log_phi_1d.broadcast_as(mu.shape())?;
        let elem = nb_log_likelihood_elem(&unseen, &mu, &log_phi)?;
        Ok((elem.mul(&scored)?.sum(1)?, scored.sum(1)?))
    }

    /// Multinomial sibling of [`Self::score_unseen_modules_nb`]: the unseen
    /// counts against `log((θβ)_m · unseen share)`, the full-axis rate
    /// restricted to what the context did not see, as the dense head does.
    pub fn score_unseen_modules_multinomial(
        &self,
        log_theta_nk: &Tensor,
        target: &ModuleTarget<'_>,
        full_km: &Tensor,
    ) -> Result<(Tensor, Tensor)> {
        let (rate_nm, unseen, share, scored) = self.unseen_parts(log_theta_nk, target, full_km)?;
        let p = rate_nm.mul(&share)?;
        let ll = (unseen * (p + 1e-20)?.log()?)?;
        Ok((ll.mul(&scored)?.sum(1)?, scored.sum(1)?))
    }

    /// `(rate [N, M], unseen counts, unseen share, scored indicator)`.
    fn unseen_parts(
        &self,
        log_theta_nk: &Tensor,
        target: &ModuleTarget<'_>,
        full_km: &Tensor,
    ) -> Result<(Tensor, Tensor, Tensor, Tensor)> {
        let rate_nm = self.mixture_rate_nd(log_theta_nk, full_km)?;
        // Both from the target row, so the difference is exact; clamp guards
        // rounding in the sums.
        let unseen = (target.values - target.visible_counts)?.clamp(0.0, f64::INFINITY)?;
        let share = target.visible_share.affine(-1.0, 1.0)?;
        let scored = share.gt(1e-6)?.to_dtype(rate_nm.dtype())?;
        Ok((rate_nm, unseen, share, scored))
    }

    /// NB log-likelihood at the sampled query genes, weighted, → `[N]`:
    /// `μ_nq = ℓ_n · (θβ)_{m(g)} · π_{g|m(g)} · exp(r_nq)`, `φ` at the module.
    /// With `r = 0` this is the expanded dictionary's rate.
    pub fn score_queries_nb(
        &self,
        log_theta_nk: &Tensor,
        q: &QueryTarget<'_>,
        full_km: &Tensor,
    ) -> Result<Tensor> {
        let rate_nm = self.mixture_rate_nd(log_theta_nk, full_km)?; // [N, M]
        let ids_m = self.coarsening.groups_of(q.gene_ids)?; // [N, Q]
        let rate_nq = rate_nm.gather(&ids_m, 1)?; // [N, Q]
        let log_factor = (self.coarsening.log_share_at(q.gene_ids)? + q.log_residual)?;
        let mu = rate_nq.mul(&log_factor.exp()?)?.broadcast_mul(q.lib)?;
        let (n, qn) = ids_m.dims2()?;
        let log_phi =
            gather_rows(&self.log_phi_1d.squeeze(0)?, &ids_m.flatten_all()?)?.reshape((n, qn))?;
        let elem = nb_log_likelihood_elem(q.values, &mu, &log_phi)?;
        elem.mul(q.weight)?.sum(1)
    }
}

/////////////////////////////////////////////////////////
// Shared scorers — one definition per likelihood       //
/////////////////////////////////////////////////////////

/// `Σ_{scored} log NB(y | residual · ℓ · rate, φ)`, summed over the last axis.
///
/// Shape-agnostic so the indexed `[N, K]` and hidden-only `[N, d_h]` heads
/// cannot drift apart: they differ in which positions they score, never in how.
/// `weight` is `None` when every position passed in is scored — the hidden-only
/// head's normal case, where there is nothing to zero out.
fn nb_score(
    values: &Tensor,
    rate: &Tensor,
    residual: Option<&Tensor>,
    lib_n1: &Tensor,
    log_phi: &Tensor,
    weight: Option<&Tensor>,
) -> Result<Tensor> {
    // μ = residual · ℓ · rate
    let mu = match residual {
        Some(r) => rate.mul(r)?.broadcast_mul(lib_n1)?,
        None => rate.broadcast_mul(lib_n1)?,
    };
    let elem = nb_log_likelihood_elem(values, &mu, log_phi)?;
    weighted_row_sum(elem, weight)
}

/// `Σ_{scored} y · log p`, summed over the last axis.
fn multinomial_score(values: &Tensor, rate: &Tensor, weight: Option<&Tensor>) -> Result<Tensor> {
    let ll = (values * (rate + 1e-20)?.log()?)?;
    weighted_row_sum(ll, weight)
}

/// Sum the last axis, weighting first where a weight was given.
fn weighted_row_sum(elem: Tensor, weight: Option<&Tensor>) -> Result<Tensor> {
    let last = elem.rank() - 1;
    match weight {
        Some(w) => elem.mul(w)?.sum(last),
        None => elem.sum(last),
    }
}

#[cfg(test)]
#[path = "masked_etm_tests.rs"]
mod masked_etm_tests;