topos 0.12.0

An autodiff compiler stack in Rust.
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
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122

//! Module checkpoints: capturing a module tree's parameter payloads
//! and restoring them into a [`Parameters`] state.
//!
//! Two identities, two tiers. The positional pair
//! ([`snapshot`]/[`restore`]) uses the module tree's stable visit
//! order — sufficient for resuming the same code, with no names
//! anywhere. The named pair ([`named_snapshot`]/[`named_restore`])
//! matches by structured [`Path`], which is what survives code
//! evolution and what foreign checkpoints (name-to-tensor maps)
//! require; missing and unexpected paths are loud errors.
//!
//! A checkpoint is pure state, so both directions are plain
//! [`Parameters`] transforms: no graph is touched, and shape
//! mismatches panic through
//! [`Parameters::with_payloads`]'s validation. The library stops at
//! the name-to-payload map; file formats stay at the edge.

use std::collections::HashMap;

use crate::{Element, Parameters, Symbol, Tensor};

use super::module::{Module, Path, named_parameters, parameters};

/// Returns the payloads of every parameter in `module`'s tree, in
/// visit order: the positional checkpoint.
///
/// # Panics
/// Panics if a visited symbol does not name a parameter `state`
/// carries.
pub fn snapshot<E: Element, M: Module<E> + ?Sized>(
    state: &Parameters<E>,
    module: &M,
) -> Vec<Tensor<E>> {
    parameters(module)
        .into_iter()
        .map(|symbol| state.of(symbol).clone())
        .collect()
}

/// Returns a new state with `module`'s parameters replaced by
/// `payloads`, matched in visit order: the positional restore.
/// Parameters outside the module keep their payloads.
///
/// # Panics
/// Panics if the payload count differs from the module's parameter
/// count, or if a payload's shape differs from its parameter's.
pub fn restore<E: Element, M: Module<E> + ?Sized>(
    state: &Parameters<E>,
    module: &M,
    payloads: Vec<Tensor<E>>,
) -> Parameters<E> {
    let symbols = parameters(module);
    assert_eq!(
        payloads.len(),
        symbols.len(),
        "the checkpoint holds {} payloads but the module has {} parameters",
        payloads.len(),
        symbols.len(),
    );
    state.with_payloads(symbols.into_iter().zip(payloads))
}

/// Returns every parameter payload in `module`'s tree with its
/// structured path, in visit order: the named checkpoint, the form
/// that survives code evolution and maps to foreign layouts.
///
/// # Panics
/// Panics as [`snapshot`] panics.
pub fn named_snapshot<E: Element, M: Module<E> + ?Sized>(
    state: &Parameters<E>,
    module: &M,
) -> Vec<(Path, Tensor<E>)> {
    named_parameters(module)
        .into_iter()
        .map(|(path, symbol)| (path, state.of(symbol).clone()))
        .collect()
}

/// Returns a new state with `module`'s parameters replaced by
/// `entries`, matched by path: the named restore.
///
/// Tied parameters (one symbol under several paths) take the last
/// matching entry in visit order. Parameters outside the module keep
/// their payloads.
///
/// # Panics
/// Panics if a module parameter has no entry, an entry matches no
/// parameter, or a payload's shape differs from its parameter's.
pub fn named_restore<E: Element, M: Module<E> + ?Sized>(
    state: &Parameters<E>,
    module: &M,
    entries: impl IntoIterator<Item = (Path, Tensor<E>)>,
) -> Parameters<E> {
    let mut entries: HashMap<Path, Tensor<E>> = entries.into_iter().collect();
    let mut replacements: Vec<(Symbol, Tensor<E>)> = Vec::new();
    let mut missing: Vec<String> = Vec::new();
    for (path, symbol) in named_parameters(module) {
        match entries.remove(&path) {
            Some(payload) => replacements.push((symbol, payload)),
            None => missing.push(path.to_string()),
        }
    }
    assert!(
        missing.is_empty(),
        "the checkpoint is missing entries for: {}",
        missing.join(", "),
    );
    assert!(
        entries.is_empty(),
        "the checkpoint holds entries no parameter matches: {}",
        entries
            .keys()
            .map(Path::to_string)
            .collect::<Vec<_>>()
            .join(", "),
    );
    state.with_payloads(replacements)
}

#[cfg(test)]
#[path = "tests/checkpoint_tests.rs"]
mod tests;