legume-numeric 0.8.11

Numeric and ML foundation for the legume ecosystem (matrix, Leiden, candle, MCMC)
Documentation 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79

use crate::candle::decoder::nb_mixture::{NbMixtureTopicDecoder, DECODER_NAME as NBMIXTURE_NAME};
use crate::candle::decoder::topic::{MultinomTopicDecoder, NbTopicDecoder};
use crate::candle::traits::model::{DecoderModuleT, EssLlikFn};
use candle_core::{Result, Tensor};
use candle_nn::VarBuilder;

/// Object-safe decoder trait for dynamic dispatch across decoder types.
///
/// All three topic decoders ignore the generic `llik` closure in
/// `DecoderModuleT::forward_with_llik` and use their own internal likelihood.
/// This trait removes the generic parameter, enabling heterogeneous collections.
pub trait DynDecoderModuleT: Send + Sync {
    fn forward(&self, z_nk: &Tensor) -> Result<Tensor>;
    fn get_dictionary(&self) -> Result<Tensor>;
    fn forward_llik(&self, z_nk: &Tensor, x_nd: &Tensor) -> Result<(Tensor, Tensor)>;
    fn dim_obs(&self) -> usize;
    fn dim_latent(&self) -> usize;
    fn build_ess_llik<'a>(&'a self, x_nd: &'a Tensor, smoothing: f64) -> Result<EssLlikFn<'a>>;
    fn decoder_name(&self) -> &str;
}

fn dummy_llik(_: &Tensor, _: &Tensor) -> Result<Tensor> {
    Err(candle_core::Error::Msg(
        "DynDecoderModuleT: external closure should never be called".into(),
    ))
}

macro_rules! impl_dyn_decoder {
    ($ty:ty, $name:expr) => {
        impl DynDecoderModuleT for $ty {
            fn forward(&self, z_nk: &Tensor) -> Result<Tensor> {
                DecoderModuleT::forward(self, z_nk)
            }
            fn get_dictionary(&self) -> Result<Tensor> {
                DecoderModuleT::get_dictionary(self)
            }
            fn forward_llik(&self, z_nk: &Tensor, x_nd: &Tensor) -> Result<(Tensor, Tensor)> {
                DecoderModuleT::forward_with_llik(self, z_nk, x_nd, &dummy_llik)
            }
            fn dim_obs(&self) -> usize {
                DecoderModuleT::dim_obs(self)
            }
            fn dim_latent(&self) -> usize {
                DecoderModuleT::dim_latent(self)
            }
            fn build_ess_llik<'a>(&'a self, x_nd: &'a Tensor, s: f64) -> Result<EssLlikFn<'a>> {
                DecoderModuleT::build_ess_llik(self, x_nd, s)
            }
            fn decoder_name(&self) -> &str {
                $name
            }
        }
    };
}

impl_dyn_decoder!(MultinomTopicDecoder, "multinom");
impl_dyn_decoder!(NbTopicDecoder, "nb");
impl_dyn_decoder!(NbMixtureTopicDecoder, NBMIXTURE_NAME);

/// Create a boxed dynamic decoder by name.
pub fn create_dyn_decoder(
    name: &str,
    n_features: usize,
    n_topics: usize,
    vs: VarBuilder,
) -> Result<Box<dyn DynDecoderModuleT>> {
    match name {
        "multinom" => Ok(Box::new(MultinomTopicDecoder::new(
            n_features, n_topics, vs,
        )?)),
        "nb" => Ok(Box::new(NbTopicDecoder::new(n_features, n_topics, vs)?)),
        NBMIXTURE_NAME => Ok(Box::new(NbMixtureTopicDecoder::new(
            n_features, n_topics, vs,
        )?)),
        other => Err(candle_core::Error::Msg(format!(
            "unknown decoder type: {other}"
        ))),
    }
}