candle_transformers/models/
with_tracing.rs

1use candle::{Module, Result, Tensor};
2use candle_nn::VarBuilder;
3
4#[derive(Debug, Clone)]
5pub struct Embedding {
6    inner: candle_nn::Embedding,
7    span: tracing::Span,
8}
9
10impl Embedding {
11    pub fn new(d1: usize, d2: usize, vb: VarBuilder) -> Result<Self> {
12        let inner = candle_nn::embedding(d1, d2, vb)?;
13        let span = tracing::span!(tracing::Level::TRACE, "embedding");
14        Ok(Self { inner, span })
15    }
16
17    pub fn from_weights(weights: Tensor) -> Result<Self> {
18        let (_in_size, out_size) = weights.dims2()?;
19        let inner = candle_nn::Embedding::new(weights, out_size);
20        let span = tracing::span!(tracing::Level::TRACE, "embedding");
21        Ok(Self { inner, span })
22    }
23
24    pub fn embeddings(&self) -> &Tensor {
25        self.inner.embeddings()
26    }
27}
28
29impl Module for Embedding {
30    fn forward(&self, xs: &Tensor) -> Result<Tensor> {
31        let _enter = self.span.enter();
32        self.inner.forward(xs)
33    }
34}
35
36#[derive(Debug, Clone)]
37pub struct Linear {
38    inner: candle_nn::Linear,
39    span: tracing::Span,
40}
41
42impl Linear {
43    pub fn from_weights(weights: Tensor, bias: Option<Tensor>) -> Self {
44        let inner = candle_nn::Linear::new(weights, bias);
45        let span = tracing::span!(tracing::Level::TRACE, "linear");
46        Self { inner, span }
47    }
48}
49
50pub fn linear_b(d1: usize, d2: usize, b: bool, vb: VarBuilder) -> Result<Linear> {
51    let inner = candle_nn::linear_b(d1, d2, b, vb)?;
52    let span = tracing::span!(tracing::Level::TRACE, "linear");
53    Ok(Linear { inner, span })
54}
55
56pub fn linear(d1: usize, d2: usize, vb: VarBuilder) -> Result<Linear> {
57    let inner = candle_nn::linear(d1, d2, vb)?;
58    let span = tracing::span!(tracing::Level::TRACE, "linear");
59    Ok(Linear { inner, span })
60}
61
62pub fn linear_no_bias(d1: usize, d2: usize, vb: VarBuilder) -> Result<Linear> {
63    let inner = candle_nn::linear_no_bias(d1, d2, vb)?;
64    let span = tracing::span!(tracing::Level::TRACE, "linear");
65    Ok(Linear { inner, span })
66}
67
68impl Module for Linear {
69    fn forward(&self, xs: &Tensor) -> Result<Tensor> {
70        let _enter = self.span.enter();
71        self.inner.forward(xs)
72    }
73}
74
75// Wrap the conv2d op to provide some tracing.
76#[derive(Debug, Clone)]
77pub struct Conv2d {
78    inner: candle_nn::Conv2d,
79    span: tracing::Span,
80}
81
82impl Module for Conv2d {
83    fn forward(&self, x: &Tensor) -> Result<Tensor> {
84        let _enter = self.span.enter();
85        self.inner.forward(x)
86    }
87}
88
89pub fn conv2d(
90    in_channels: usize,
91    out_channels: usize,
92    kernel_size: usize,
93    cfg: candle_nn::Conv2dConfig,
94    vs: candle_nn::VarBuilder,
95) -> Result<Conv2d> {
96    let span = tracing::span!(tracing::Level::TRACE, "conv2d");
97    let inner = candle_nn::conv2d(in_channels, out_channels, kernel_size, cfg, vs)?;
98    Ok(Conv2d { inner, span })
99}
100
101// QMatMul wrapper adding some tracing.
102#[derive(Clone)]
103pub struct QMatMul {
104    inner: candle::quantized::QMatMul,
105    span: tracing::Span,
106}
107
108impl QMatMul {
109    pub fn new(
110        out_dim: usize,
111        in_dim: usize,
112        vb: crate::quantized_var_builder::VarBuilder,
113    ) -> Result<Self> {
114        let ws = vb.get((in_dim, out_dim), "weight")?;
115        let inner = candle::quantized::QMatMul::from_arc(ws)?;
116        let span = tracing::span!(tracing::Level::TRACE, "qmatmul");
117        Ok(Self { inner, span })
118    }
119
120    pub fn from_weights(ws: std::sync::Arc<candle::quantized::QTensor>) -> Result<Self> {
121        let inner = candle::quantized::QMatMul::from_arc(ws)?;
122        let span = tracing::span!(tracing::Level::TRACE, "qmatmul");
123        Ok(Self { inner, span })
124    }
125}
126
127impl Module for QMatMul {
128    fn forward(&self, xs: &Tensor) -> Result<Tensor> {
129        let _enter = self.span.enter();
130        self.inner.forward(xs)
131    }
132}
133
134impl std::fmt::Debug for QMatMul {
135    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
136        write!(f, "QMatMul")
137    }
138}
139
140#[derive(Clone, Debug)]
141pub struct LayerNorm {
142    inner: candle_nn::LayerNorm,
143    span: tracing::Span,
144}
145
146impl LayerNorm {
147    pub fn new(weight: Tensor, bias: Tensor, eps: f64) -> Self {
148        let inner = candle_nn::LayerNorm::new(weight, bias, eps);
149        let span = tracing::span!(tracing::Level::TRACE, "layer-norm");
150        Self { inner, span }
151    }
152}
153
154impl Module for LayerNorm {
155    fn forward(&self, xs: &Tensor) -> Result<Tensor> {
156        let _enter = self.span.enter();
157        self.inner.forward(xs)
158    }
159}
160
161pub fn layer_norm<C: Into<candle_nn::LayerNormConfig>>(
162    size: usize,
163    c: C,
164    vb: VarBuilder,
165) -> Result<LayerNorm> {
166    let inner = candle_nn::layer_norm(size, c, vb)?;
167    let span = tracing::span!(tracing::Level::TRACE, "layer-norm");
168    Ok(LayerNorm { inner, span })
169}
170
171#[derive(Debug, Clone)]
172pub struct RmsNorm {
173    inner: candle_nn::RmsNorm,
174    span: tracing::Span,
175}
176
177impl RmsNorm {
178    pub fn new(size: usize, eps: f64, vb: VarBuilder) -> Result<Self> {
179        let span = tracing::span!(tracing::Level::TRACE, "rms-norm");
180        let inner = candle_nn::rms_norm(size, eps, vb)?;
181        Ok(Self { inner, span })
182    }
183
184    pub fn forward_diff(&self, x: &Tensor) -> Result<Tensor> {
185        let _enter = self.span.enter();
186        self.inner.forward_diff(x)
187    }
188}
189
190impl Module for RmsNorm {
191    fn forward(&self, x: &Tensor) -> Result<Tensor> {
192        let _enter = self.span.enter();
193        self.inner.forward(x)
194    }
195}
candle_transformers/models/with_tracing.rs

candle_transformers/models/
with_tracing.rs
