metrale_circuit/checkpoint/
states.rs

1// SPDX-License-Identifier: MIT OR Apache-2.0
2
3//! 2026-09-30: The recurrent state an embedded circuit declares, for a model the engine has
4//! already configured: the circuit is chosen by `model_type` and the declarations are
5//! evaluated under dims the caller reads from its own (tensor-parallel local) config, so no
6//! precision and no config.json is involved.
7//!
8//! Owner: metrale-circuit.
9//! Invariants:
10//! - `Ok(None)` only when no embedded circuit serves `model_type`; a circuit that serves it
11//!   but cannot be read, or whose layers hold recurrent state in more than one block, is an
12//!   error.
13
14use std::collections::BTreeMap;
15
16use super::{ARCHES, BLOCKS, config_maps};
17use crate::circuit_toml::{CircuitError, parse_file};
18use crate::instantiate::state_decl;
19use crate::ir::Section;
20use crate::state::{StateDecl, StateKind};
21
22/// 2026-09-30: The recurrent states of one layer of the circuit that serves the engine-configured
23/// `model_type` (`ConfigMap::serves_engine_model_type`), under
24/// `dims`: the declarations of the one layer block that keeps recurrent state (GatedDeltaNet
25/// `gdn`, Mamba2 `mamba`), ids `<block>.<state>`. `Ok(Some(empty))` for a circuit with no
26/// recurrent layer.
27pub fn recurrent_states(
28    model_type: &str,
29    dims: &BTreeMap<String, u64>,
30) -> Result<Option<Vec<StateDecl>>, CircuitError> {
31    let maps = config_maps().map_err(|e| CircuitError::Parse(e.to_string()))?;
32    let Some(i) = maps
33        .iter()
34        .position(|m| m.serves_engine_model_type(model_type))
35    else {
36        return Ok(None);
37    };
38    let file = parse_file(ARCHES[i].circuit, &BLOCKS)?;
39    let mut templates: Vec<&String> = file
40        .layout
41        .blocks
42        .values()
43        .flatten()
44        .filter(|t| {
45            file.block
46                .get(*t)
47                .is_some_and(|b| b.state.iter().any(|s| s.kind == "recurrent"))
48        })
49        .collect();
50    templates.sort();
51    templates.dedup();
52    let template = match templates.as_slice() {
53        [] => return Ok(Some(Vec::new())),
54        [one] => *one,
55        many => {
56            return Err(CircuitError::Layout(format!(
57                "circuit `{}`: recurrent state in more than one layer block ({many:?})",
58                file.arch
59            )));
60        }
61    };
62    let mut out = Vec::new();
63    let block = &file.block[template].state;
64    for sf in block {
65        let d = state_decl(template, template, None, Section::Main, sf, block, dims)?;
66        if d.kind == StateKind::Recurrent {
67            out.push(d);
68        }
69    }
70    Ok(Some(out))
71}
metrale_circuit/checkpoint/states.rs

metrale_circuit/checkpoint/
states.rs
