candle_graph/
known.rs

1//! Ground truth for the candle-nn constructors that register parameters.
2//!
3//! Every entry here was read out of candle-nn 0.11.0 source rather than recalled, because the
4//! whole output's correctness rests on these names matching what `VarBuilder::get` actually
5//! stores. Citations are to
6//! `~/.cargo/registry/src/*/candle-nn-0.11.0/src/`.
7
8use serde::Serialize;
9
10/// What a leaf tensor is for. Kept coarse on purpose: it drives display grouping and the
11/// "is this trainable" question, not arithmetic.
12#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize)]
13pub enum ParamKind {
14    Weight,
15    Bias,
16    RunningMean,
17    RunningVar,
18    /// Registered through a raw `vb.get(..)` / `vb.get_with_hints(..)` call in user code.
19    Raw,
20}
21
22/// One leaf tensor a constructor registers, relative to the `VarBuilder` it was handed.
23#[derive(Debug, Clone, Copy)]
24pub struct Leaf {
25    pub name: &'static str,
26    pub kind: ParamKind,
27    /// False when the tensor is only registered for some configurations (e.g. `layer_norm`'s
28    /// bias, which exists only when `affine`). Callers must mark these conditional unless they
29    /// can resolve the config.
30    pub unconditional: bool,
31    /// True when the tensor *name* is formatted from the constructor's config argument rather
32    /// than being a fixed literal. `LSTM::new` builds `weight_ih_l{layer_idx}{direction}`
33    /// (rnn.rs:139-147), so [`Leaf::name`] is only the correct key for the default config.
34    /// Callers that cannot resolve the config must emit a name family, not this literal.
35    pub config_named: bool,
36}
37
38const fn leaf(name: &'static str, kind: ParamKind) -> Leaf {
39    Leaf {
40        name,
41        kind,
42        unconditional: true,
43        config_named: false,
44    }
45}
46
47const fn cond(name: &'static str, kind: ParamKind) -> Leaf {
48    Leaf {
49        name,
50        kind,
51        unconditional: false,
52        config_named: false,
53    }
54}
55
56/// A leaf whose name is derived from the config argument; `name` is the default-config spelling.
57const fn named_by_config(name: &'static str, kind: ParamKind, unconditional: bool) -> Leaf {
58    Leaf {
59        name,
60        kind,
61        unconditional,
62        config_named: true,
63    }
64}
65
66/// A candle-nn constructor whose parameter layout we know exactly.
67#[derive(Debug, Clone, Copy)]
68pub struct Constructor {
69    /// Last path segment of the call, e.g. `linear` for `candle_nn::linear` or `nn::linear`.
70    pub func: &'static str,
71    /// Zero-based index of the `VarBuilder` argument.
72    pub vb_arg: usize,
73    pub leaves: &'static [Leaf],
74    /// Source citation, carried into diagnostics so a reader can check us.
75    pub cite: &'static str,
76}
77
78/// `linear` and friends: `weight` always, `bias` only in the biased variants.
79/// linear.rs:84-95 (`linear`), :97-101 (`linear_no_bias`), :103 (`linear_b`).
80const LINEAR: &[Leaf] = &[
81    leaf("weight", ParamKind::Weight),
82    leaf("bias", ParamKind::Bias),
83];
84const LINEAR_NO_BIAS: &[Leaf] = &[leaf("weight", ParamKind::Weight)];
85/// `linear_b(.., bias: bool, ..)` — the bias is config-dependent.
86const LINEAR_B: &[Leaf] = &[
87    leaf("weight", ParamKind::Weight),
88    cond("bias", ParamKind::Bias),
89];
90
91/// embedding.rs:39-49 — the tensor is named "weight" despite the binding being `embeddings`.
92const EMBEDDING: &[Leaf] = &[leaf("weight", ParamKind::Weight)];
93
94/// layer_norm.rs:146-163. `weight` is unconditional; `bias` exists only when `config.affine`.
95/// Note `impl From<f64> for LayerNormConfig` (:52-59) sets `affine: true`, so the common
96/// `layer_norm(dim, 1e-5, vb)` form does register a bias.
97const LAYER_NORM: &[Leaf] = &[
98    leaf("weight", ParamKind::Weight),
99    cond("bias", ParamKind::Bias),
100];
101/// layer_norm.rs:166 — explicitly no bias.
102const LAYER_NORM_NO_BIAS: &[Leaf] = &[leaf("weight", ParamKind::Weight)];
103/// layer_norm.rs:212-219 builds a `LayerNormConfig { affine: false, .. }`, so weight only.
104const RMS_NORM: &[Leaf] = &[leaf("weight", ParamKind::Weight)];
105
106/// group_norm.rs:76-84 — affine weight and bias are always registered.
107const GROUP_NORM: &[Leaf] = &[
108    leaf("weight", ParamKind::Weight),
109    leaf("bias", ParamKind::Bias),
110];
111
112/// activation.rs:104-108 — PReLU always stores one scalar or one value per channel as `weight`.
113const PRELU: &[Leaf] = &[leaf("weight", ParamKind::Weight)];
114
115/// conv.rs:307-327 (`conv1d`), :386-410 (`conv2d`), and the transpose/no-bias variants.
116const CONV: &[Leaf] = &[
117    leaf("weight", ParamKind::Weight),
118    leaf("bias", ParamKind::Bias),
119];
120const CONV_NO_BIAS: &[Leaf] = &[leaf("weight", ParamKind::Weight)];
121
122/// rnn.rs:304-343 (`GRU::new`). Names are hard-coded `_l0` literals — `GRUConfig` has no
123/// layer/direction fields (rnn.rs:257-262) — so the keys are always exactly these. Both biases
124/// are conditional: `GRUConfig::default()` (rnn.rs:264) supplies them, `default_no_bias()`
125/// (rnn.rs:276) sets both inits to `None` and registers nothing.
126const GRU: &[Leaf] = &[
127    leaf("weight_ih_l0", ParamKind::Weight),
128    leaf("weight_hh_l0", ParamKind::Weight),
129    cond("bias_ih_l0", ParamKind::Bias),
130    cond("bias_hh_l0", ParamKind::Bias),
131];
132
133/// rnn.rs:134-187 (`LSTM::new`). Unlike GRU, every name is built with `format!` from the config:
134/// `weight_ih_l{layer_idx}{direction}` (rnn.rs:145-147), where `direction` is `""` for
135/// `Direction::Forward` and `"_reverse"` for `Backward` (rnn.rs:141-144). The literals below are
136/// therefore only the *default-config* spelling and are marked `config_named` so an unresolved
137/// config produces a name family instead of a wrong key.
138const LSTM: &[Leaf] = &[
139    named_by_config("weight_ih_l0", ParamKind::Weight, true),
140    named_by_config("weight_hh_l0", ParamKind::Weight, true),
141    named_by_config("bias_ih_l0", ParamKind::Bias, false),
142    named_by_config("bias_hh_l0", ParamKind::Bias, false),
143];
144
145/// batch_norm.rs:301-317. Running stats are always registered; affine params are conditional.
146/// These tensors are commonly excluded from optimizer variable lists, so flagging them
147/// distinctly is worth the extra kinds.
148const BATCH_NORM: &[Leaf] = &[
149    leaf("running_mean", ParamKind::RunningMean),
150    leaf("running_var", ParamKind::RunningVar),
151    cond("weight", ParamKind::Weight),
152    cond("bias", ParamKind::Bias),
153];
154
155pub const CONSTRUCTORS: &[Constructor] = &[
156    Constructor {
157        func: "linear",
158        vb_arg: 2,
159        leaves: LINEAR,
160        cite: "linear.rs:84",
161    },
162    Constructor {
163        func: "linear_no_bias",
164        vb_arg: 2,
165        leaves: LINEAR_NO_BIAS,
166        cite: "linear.rs:97",
167    },
168    Constructor {
169        func: "linear_b",
170        vb_arg: 3,
171        leaves: LINEAR_B,
172        cite: "linear.rs:103",
173    },
174    Constructor {
175        func: "embedding",
176        vb_arg: 2,
177        leaves: EMBEDDING,
178        cite: "embedding.rs:39",
179    },
180    Constructor {
181        func: "layer_norm",
182        vb_arg: 2,
183        leaves: LAYER_NORM,
184        cite: "layer_norm.rs:146",
185    },
186    Constructor {
187        func: "layer_norm_no_bias",
188        vb_arg: 2,
189        leaves: LAYER_NORM_NO_BIAS,
190        cite: "layer_norm.rs:166",
191    },
192    Constructor {
193        func: "rms_norm",
194        vb_arg: 2,
195        leaves: RMS_NORM,
196        cite: "layer_norm.rs:212",
197    },
198    Constructor {
199        func: "group_norm",
200        vb_arg: 3,
201        leaves: GROUP_NORM,
202        cite: "group_norm.rs:76",
203    },
204    Constructor {
205        func: "prelu",
206        vb_arg: 1,
207        leaves: PRELU,
208        cite: "activation.rs:104",
209    },
210    Constructor {
211        func: "conv1d",
212        vb_arg: 4,
213        leaves: CONV,
214        cite: "conv.rs:307",
215    },
216    Constructor {
217        func: "conv1d_no_bias",
218        vb_arg: 4,
219        leaves: CONV_NO_BIAS,
220        cite: "conv.rs:329",
221    },
222    Constructor {
223        func: "conv2d",
224        vb_arg: 4,
225        leaves: CONV,
226        cite: "conv.rs:386",
227    },
228    Constructor {
229        func: "conv2d_no_bias",
230        vb_arg: 4,
231        leaves: CONV_NO_BIAS,
232        cite: "conv.rs:413",
233    },
234    Constructor {
235        func: "conv_transpose1d",
236        vb_arg: 4,
237        leaves: CONV,
238        cite: "conv.rs:345",
239    },
240    Constructor {
241        func: "conv_transpose1d_no_bias",
242        vb_arg: 4,
243        leaves: CONV_NO_BIAS,
244        cite: "conv.rs:366",
245    },
246    Constructor {
247        func: "conv_transpose2d",
248        vb_arg: 4,
249        leaves: CONV,
250        cite: "conv.rs:434",
251    },
252    Constructor {
253        func: "conv_transpose2d_no_bias",
254        vb_arg: 4,
255        leaves: CONV_NO_BIAS,
256        cite: "conv.rs:455",
257    },
258    Constructor {
259        func: "batch_norm",
260        vb_arg: 2,
261        leaves: BATCH_NORM,
262        cite: "batch_norm.rs:301",
263    },
264    Constructor {
265        func: "gru",
266        vb_arg: 3,
267        leaves: GRU,
268        cite: "rnn.rs:345",
269    },
270    Constructor {
271        func: "lstm",
272        vb_arg: 3,
273        leaves: LSTM,
274        cite: "rnn.rs:189",
275    },
276];
277
278pub fn lookup(func: &str) -> Option<&'static Constructor> {
279    CONSTRUCTORS.iter().find(|c| c.func == func)
280}
281
282/// `VarBuilder` methods that register a tensor directly, with the index of the name argument.
283/// var_builder.rs:203 (`get_with_hints`), :213 (`get`), :218 (`get_unchecked`).
284pub fn raw_get_name_arg(method: &str) -> Option<usize> {
285    match method {
286        "get" => Some(1),
287        "get_with_hints" => Some(1),
288        "get_with_hints_dtype" => Some(1),
289        "get_unchecked" => Some(0),
290        "get_unchecked_dtype" => Some(0),
291        _ => None,
292    }
293}
294
295/// `VarBuilder` methods that return a *re-prefixed* builder rather than a tensor.
296/// var_builder.rs:162 (`pp`), :150 (`push_prefix`), :139 (`set_prefix`), :129 (`root`).
297pub fn prefix_method(method: &str) -> Option<PrefixOp> {
298    match method {
299        "pp" | "push_prefix" => Some(PrefixOp::Push),
300        "set_prefix" => Some(PrefixOp::Replace),
301        "root" => Some(PrefixOp::Root),
302        _ => None,
303    }
304}
305
306#[derive(Debug, Clone, Copy, PartialEq, Eq)]
307pub enum PrefixOp {
308    Push,
309    Replace,
310    Root,
311}
candle_graph/known.rs

candle_graph/
known.rs
